魔改labelimg做出的数据标注小工具

 1.前言：

作业的要求是写一个数据标注的小工具，可以参考labelimg，参考精灵标注，但我说实话，这不是一个上了6节pyqt就能搞出来的东西，你黄哥也是费尽千幸万苦，最终不断地加功能，不断的改bug，终于实现了老师心中一博昏的作品。

因为是帮助广大（每一年的大作业都是这个）学生，所以写下这篇文章，(自我感动.jpg)

太感动了，太有实力了你黄哥。

好吧，接下来欣赏我和chatgpt的热血组合技吧（说实话，chatgpt太笨了，功能都得自己加）

2.界面演示

也算是labelimg有的功能基本都有了，和精灵标注相比，更是一个天上一个地下了（差不多得了）

实际上还是有个bug的后面会讲，其实就是你必须点开一个label才能运行，要不然会有bug

这个的流程就是，选择一张图片，然后跳出让你选择一个label（主要是针对大项目的打标签做训练集，这个label都是会给你的，但其实也能自己加的）然后可以选择颜色，要移动选框就点击开关，还可以选择上一张下一张，还有保存xml，json，txt各种形式...

2.结构说明：

分为了三部分，画框（frame），标注（Marking），和各种按钮的交互（Data annotation）

要运行Data annotation，即可，代码注释我也已经有了

emmm，这里直接上代码ok？

3.frame：

from PyQt5 import QtCore, QtGui, QtWidgets
from PyQt5.QtWidgets import QWidget, QApplication, QLabel, QMessageBox, QPushButton
from PyQt5.QtCore import QRect, Qt
from PyQt5.QtGui import QPixmap, QPainter, QPen
from Marking import markinguse
import sys
import qtmodern.styles
import qtmodern.windows
#定义QLabel实现各种鼠标操作
class Mylabel(QLabel):
    zgk = False
    baglist=[]
    def __init__(self,parent=None):
        #初始化参数
        super(Mylabel,self).__init__(parent)
        self.initParam()
    #初始化值
    def initParam(self):
        self.x_first=0
        self.y_first=0
        self.x_second=0
        self.y_second=0
        self.xtrue=0
        self.ytrue=0
        self.rect=QRect()
        self.flag=False#标志
        #定义一个baglist，用来储存框
        self.labelindex=0
        self.curchoosebag=[]#删除的bag
        self.curlabelindex=-1
        self.deletebagflag=False
        self.fileInfo={}
        self.color=None
    #定义画框工具，这里基本是QPen的
    #单击画矩形框
    def mousePressEvent(self, QMouseEvent):
        x = QMouseEvent.pos().x()
        y = QMouseEvent.pos().y()

        self.curchoosebag = []
        self.deletebagflag = False
        for index, bag in enumerate(Mylabel.baglist):
            if bag[0] <= x <= bag[2] and bag[1] <= y <= bag[3]and Mylabel.zgk==True:
                # 检查是否点击了已存在的框
                self.curchoosebag = bag
                self.curlabelindex = index
                self.flag = True
                self.update()
                return  # 如果点击在框内，就不创建新框

        # 如果没有点击在任何框内，则创建新框
        self.flag = True
        self.x_first = x
        self.y_first = y
    #鼠标移动的过程
    def mouseMoveEvent(self, QMouseEvent):
        if self.flag:
            if self.curchoosebag:
                # 计算移动距离
                self.cs = 1
                dx = QMouseEvent.pos().x() - self.x_first
                dy = QMouseEvent.pos().y() - self.y_first

                # 更新框的位置
                self.curchoosebag = (self.curchoosebag[0] + dx, self.curchoosebag[1] + dy,
                                     self.curchoosebag[2] + dx, self.curchoosebag[3] + dy, self.curchoosebag[4],self.curchoosebag[5],self.curchoosebag[6])

                # 更新列表中的框位置
                Mylabel.baglist[self.curlabelindex]=self.curchoosebag





                # 为下一次移动更新初始坐标
                self.x_first = QMouseEvent.pos().x()
                self.y_first = QMouseEvent.pos().y()
            else:
                # 如果不是移动框，则正常绘制新框
                self.xtrue = QMouseEvent.pos().x()
                self.ytrue = QMouseEvent.pos().y()
            self.update()

    def mouseReleaseEvent(self, QMouseEvent):
        if self.flag and not self.curchoosebag:
            # 如果是创建新框的行为，完成创建
            self.x_second = QMouseEvent.pos().x()
            self.y_second = QMouseEvent.pos().y()

            #实时操作流畅准确
            self.xtrue=self.x_first
            self.ytrue=self.y_first
            #设置误触事件，如果很小的移动或者指点了一下，就不画框了
            if self.x_first==self.x_second or self.y_second==self.y_first:
                return
            #储存数据到bag里面
            #这里先要标注，使用Marking.py这个文件
            marke=markinguse()
            if marke.exec_():
                labelname = marke.getValue()
                self.savebaglist(self.x_first, self.y_first, self.x_second, self.y_second,labelname,self.color)
                print(Mylabel.baglist)
             # 无论是移动框还是创建新框，都要重置状态
        self.flag = False
        self.curchoosebag = []
        self.curlabelindex = -1

        self.update()

        QMouseEvent.ignore()


    #使用QPen进行绘制
    def paintEvent(self, QPaintEvent):
        super().paintEvent(QPaintEvent)
        painter=QPainter()
        painter.begin(self)

        #这个是最终的框
        for point in Mylabel.baglist:
            rect = QRect(point[0], point[1], abs(point[0] - point[2]), abs(point[1] - point[3]))
            painter.setPen(QPen(point[6], 2, Qt.SolidLine))
            painter.drawRect(rect)
            painter.drawText(point[0], point[1], point[4])
        #这个是最开始的实时框
        if not self.curchoosebag and self.flag:
            tempx_first = min(self.x_first, self.xtrue)
            tempy_first = min(self.y_first, self.ytrue)
            tempx_second = max(self.x_first, self.xtrue)
            tempy_second = max(self.y_first, self.ytrue)
            width = tempx_second - tempx_first
            height = tempy_second - tempy_first
            currect = QRect(tempx_first, tempy_first, width, height)
            painter.setPen(QPen(Qt.blue, 1, Qt.SolidLine))
            painter.drawRect(currect)




        painter.end()

    #鼠标双击，可以删除框，框的使用和改变，感觉这个有点难以描述
    def mouseDoubleClickEvent(self, QMouseEvent):
        x=QMouseEvent.pos().x()
        y=QMouseEvent.pos().y()
        self.curchoosebag=[]
        #如果已经没了就直接return了
        if Mylabel.baglist==[]:
            return
        else:
            tempbaglist=Mylabel.baglist
            for index,bag in enumerate(tempbaglist):
                if bag[0]<=x<=bag[2] and bag[1]<=y<=bag[3]:
                    #表示鼠标的实时位置在框内，这里可以加内容感觉
                    self.curchoosebag=bag
                    self.curlabelindex=index
                    self.update()
                if self.curchoosebag!=[]:
                    #设置选择按钮
                    reply = QMessageBox.question(self, "警告！", "是否要删除当前选中的标注框",
                                                 QMessageBox.StandardButton.Yes | QMessageBox.StandardButton.No,
                                                 QMessageBox.StandardButton.Yes)

                    if reply == QMessageBox.Yes:
                        self.deletebagflag = True
                        Mylabel.baglist.pop(self.curlabelindex)
                        self.update()
                    else:

                        return

    #保存baglist
    def savebaglist(self,xfirst,yfirst,xsecond,ysecond,labelname,color):
        tempx_first=min(xfirst,xsecond)
        tempx_second=max(xsecond,xfirst)
        tempy_first=min(yfirst,ysecond)
        tempy_second=max(ysecond,yfirst)
        marke = markinguse()  # 在循环内部创建新的markinguse对象
        self.color=marke.color
        bag=(tempx_first,tempy_first,tempx_second,tempy_second,labelname,self.labelindex,self.color)
        Mylabel.baglist.append(bag)
        self.labelindex+=1

    def resetBoxes(self):
        Mylabel.baglist = []  # 假设这是存储框信息的列表
        self.update()  # 更新界面，重绘标签
    def openzgk(self):
        Mylabel.zgk=not Mylabel.zgk
        print(Mylabel.zgk)

4.Marking:

#此代码的用处是写label
from PyQt5 import QtCore, QtGui, QtWidgets
import qtmodern.styles
import qtmodern.windows
from PyQt5.QtWidgets import QFileDialog
class marking(object):
    #设置添加标签部分的控件
    def __init__(self):
        super().__init__()
        self.setWindowTitle("添加标签")

    def setupUi(self, Dialog):
        Dialog.setObjectName("Dialog")
        #适当调整高度
        Dialog.resize(285, 336)
        sizePolicy = QtWidgets.QSizePolicy(QtWidgets.QSizePolicy.Preferred, QtWidgets.QSizePolicy.Preferred)
        sizePolicy.setHorizontalStretch(0)
        sizePolicy.setVerticalStretch(0)
        sizePolicy.setHeightForWidth(Dialog.sizePolicy().hasHeightForWidth())
        Dialog.setSizePolicy(sizePolicy)
        Dialog.setMinimumSize(QtCore.QSize(285, 336))
        Dialog.setMaximumSize(QtCore.QSize(285, 336))
        #设置一个按钮


        self.buttonBox = QtWidgets.QDialogButtonBox(Dialog)
        self.buttonBox.setGeometry(QtCore.QRect(80, 39, 193, 28))#大小
        self.buttonBox.setOrientation(QtCore.Qt.Horizontal)#方向
        self.buttonBox.setStandardButtons(QtWidgets.QDialogButtonBox.Cancel | QtWidgets.QDialogButtonBox.Ok)#ok||cancel
        self.buttonBox.setObjectName("buttonBox")
        #这两个都是无法改动的


        self.leditChoosedLabel = QtWidgets.QLineEdit(Dialog)
        self.leditChoosedLabel.setGeometry(QtCore.QRect(11, 11, 261, 21))
        self.leditChoosedLabel.setObjectName("leditChoosedLabel")
        self.leditChoosedLabel.setEnabled(True)
        self.lviewLabelList = QtWidgets.QListView(Dialog)
        self.lviewLabelList.setGeometry(QtCore.QRect(10, 80, 261, 241))
        self.lviewLabelList.setEditTriggers(QtWidgets.QAbstractItemView.NoEditTriggers)
        self.lviewLabelList.setObjectName("lviewLabelList")
        QtCore.QMetaObject.connectSlotsByName(Dialog)



#正式开始编写按钮
from PyQt5 import QtCore, QtGui, QtWidgets
from PyQt5.QtWidgets import QColorDialog,QMainWindow, QApplication,  QDialog,QMessageBox
from PyQt5.QtCore import QStringListModel

class markinguse(QDialog, marking):
    labelList = []  # 类级别的变量，用于保存文件内容
    color=None

    def __init__(self, parent=None):
        QDialog.__init__(self, parent)
        self.setupUi(self)
        self.initLableList()
        self.lviewLabelList.clicked.connect(self.clickedlist)
        self.buttonBox.accepted.connect(self.validate)
        self.buttonBox.rejected.connect(self.reject)


    def initLableList(self):
        if not markinguse.labelList:  # 只在 labelList 为空时执行
            filename, _ = QFileDialog.getOpenFileName(self, "Open Label File", "", "Text Files (*.txt)")
            if filename:
                with open(filename, 'r', encoding='utf-8') as f:
                    markinguse.labelList = [line.strip() for line in f]
        self.labelslm = QStringListModel()
        self.labelslm.setStringList(markinguse.labelList)
        self.lviewLabelList.setModel(self.labelslm)

    def clickedlist(self, qModelIndex):
        self.leditChoosedLabel.setText(markinguse.labelList[qModelIndex.row()])

    def getValue(self):

        if self.leditChoosedLabel.text() != '':
            return self.leditChoosedLabel.text()
        elif self.lviewLabelList.currentIndex().row() != -1:
            return markinguse.labelList[self.lviewLabelList.currentIndex().row()]
        else:
            return ''

    def addCustomLabel(self, custom_label):
        if custom_label not in markinguse.labelList:
            markinguse.labelList.append(custom_label)  # 将用户输入的标签添加到类级别的标签列表中
            self.labelslm.setStringList(markinguse.labelList)  # 更新界面的标签列表
    def validate(self):
        if self.leditChoosedLabel.text() != '' or self.lviewLabelList.currentIndex().row() != -1:
            custom_label = self.getValue()  # 获取用户输入的标签内容
            self.addCustomLabel(custom_label)  # 将用户输入的标签添加到标签列表中
            self.accept()
        else:
            QMessageBox.warning(self, '警告', '请选择一个标签或输入一个新标签！')

    def clickedlist(self, qModelIndex):
        self.leditChoosedLabel.setText(markinguse.labelList[qModelIndex.row()])
        # 弹出颜色选择对话框
        markinguse.color = QColorDialog.getColor()

    def clearData(self):
        markinguse.labelList=[]
        self.labelslm.setStringList([])  # 清空标签列表
        self.labelList = []  # 清空标签内容
        self.leditChoosedLabel.setText('')  # 清空选择的标签文本


if __name__ == "__main__":
    import sys

    app = QtWidgets.QApplication(sys.argv)
    mark = markinguse()
    print('dialogChooseLabel.getValue()=', mark.getValue())
    sys.exit(app.exec_())

5.Data annotation

import qtmodern.styles
import qtmodern.windows
from PyQt5 import QtCore, QtGui, QtWidgets
from PyQt5.QtWidgets import QApplication, QMainWindow, QWidget, QFileDialog, QScrollArea, QVBoxLayout,QPushButton,QTextEdit
from PyQt5.QtGui import QPixmap, QPainter, QPen
from PyQt5.QtCore import QRect, Qt, QDir
from frame import Mylabel
import sys, os
from Marking import markinguse
import xml.etree.ElementTree as ET
import json
#设计一个导入文件和保存文件的代码
class Ui_Form(object):

    def setupUi(self, Form):
        Form.setObjectName("Form")
        Form.setWindowTitle("图片标注小工具")
        Form.resize(1000, 800)
        self.layoutWidget = QtWidgets.QWidget(Form)
        self.layoutWidget.setGeometry(QtCore.QRect(1, 1, 1000, 600))
        self.layoutWidget.setObjectName("layoutWidget")
        self.verticalLayout = QtWidgets.QVBoxLayout(self.layoutWidget)
        self.verticalLayout.setContentsMargins(0, 0, 0, 0)
        self.verticalLayout.setObjectName("verticalLayout")
        self.label = Mylabel(self.layoutWidget)
        self.label.resize(1200, 600)
        self.label.setObjectName("label")
        # 添加居中展示
        self.label.setAlignment(Qt.AlignCenter)

        self.verticalLayout.addWidget(self.label)

        # 添加滚动栏
        self.scroll_area = QScrollArea()
        self.scroll_area.setWidget(self.label)
        self.scroll_area.setWidgetResizable(True)
        self.verticalLayout.addWidget(self.scroll_area)

        self.pushButton = QtWidgets.QPushButton(self.layoutWidget)
        self.pushButton.setObjectName("pushButton")
        self.verticalLayout.addWidget(self.pushButton)

        self.pushButtonopen = QtWidgets.QPushButton(self.layoutWidget)
        self.pushButtonopen.setObjectName("pushButtonsave")
        self.verticalLayout.addWidget(self.pushButtonopen)

        self.pushButtonopenonly = QtWidgets.QPushButton(self.layoutWidget)
        self.pushButtonopenonly.setObjectName("pushButtonsaveonly")
        self.verticalLayout.addWidget(self.pushButtonopenonly)

        self.pushButtonclear = QtWidgets.QPushButton(self.layoutWidget)
        self.pushButtonclear.setObjectName("pushButtonclear")
        self.verticalLayout.addWidget(self.pushButtonclear)

        self.pushButtonzgk = QtWidgets.QPushButton(self.layoutWidget)
        self.pushButtonzgk.setObjectName("pushButtonzgk")
        self.verticalLayout.addWidget(self.pushButtonzgk)



        self.retranslateUi(Form)

    def retranslateUi(self, Form):
        _translate = QtCore.QCoreApplication.translate
        Form.setWindowTitle(_translate("Form", "Form"))
        self.label.setText(_translate("Form", "请导入图片"))
        self.pushButton.setText(_translate("Form", "保存"))
        self.pushButtonopen.setText(_translate("Form", "打开文件"))
        self.pushButtonopenonly.setText(_translate("Form","打开一张图片"))
        self.pushButtonclear.setText(_translate("Form", "清空并重新选择label"))
        self.pushButtonzgk.setText(_translate("Form", "打开移动图像开关,再按一次关闭"))


class MyMainWindow(QWidget, Ui_Form):
    def __init__(self, parent=None):
        super(MyMainWindow, self).__init__(parent)
        self.setupUi(self)
        self.initUI()
        self.fileName = ""  # 用于存储当前显示的图片文件名
        self.img = None  # 用于存储当前显示的图片
        self.MylabeL = Mylabel()
        self.imageList=[]
        self.currentIndex = 0  # 用于存储当前显示的图片在列表中的索引
        self.saveboolen=1
        self.pushButton.clicked.connect(self.onSave)
        self.pushButtonopen.clicked.connect(self.onOpen)
        self.pushButtonopenonly.clicked.connect(self.onOpenonly)
        self.pushButtonclear.clicked.connect(self.clearlabel)
        self.pushButtonzgk.clicked.connect(self.opnezgk)
        self.pushButtonPrev.clicked.connect(self.onPrevImage)  # 连接上一张按钮
        self.pushButtonNext.clicked.connect(self.onNextImage)  # 连接下一张按钮

    def initUI(self):
        # 新增的按钮
        self.pushButtonPrev = QtWidgets.QPushButton(self.layoutWidget)
        self.pushButtonPrev.setObjectName("pushButtonPrev")
        self.pushButtonPrev.setText("上一张")
        self.verticalLayout.addWidget(self.pushButtonPrev)

        self.pushButtonNext = QtWidgets.QPushButton(self.layoutWidget)
        self.pushButtonNext.setObjectName("pushButtonNext")
        self.pushButtonNext.setText("下一张")
        self.verticalLayout.addWidget(self.pushButtonNext)



    def onSave(self):
        curPath = QDir.currentPath()  # 获取系统当前目录
        title = "保存标注文件格式"
        filt = "Text Format (*.txt);;Json Format(*.json);;XML Format(*.XML)"
        saveFileName, flt = QFileDialog.getSaveFileName(self, title, curPath, filt)
        import os
        if saveFileName != '':
            fileName, suffixName = os.path.splitext(os.path.basename(saveFileName))
            if suffixName == ".txt":
                self.savetoText(saveFileName,self.MylabeL.baglist)
            elif suffixName == ".csv":
                self.savetoCSV(saveFileName)
            elif suffixName == ".json":
                self.savetoJson(saveFileName,self.MylabeL.baglist)
            elif suffixName == ".XML":
                self.savetoXML(saveFileName,self.MylabeL.baglist)
            else:
                pass
        else:
            return

    def savetoText(self, fileName, baglist):
        with open(fileName, 'w') as file:
            for item in baglist:
                line = '{} {} {} {} {} {} {}\n'.format(item[0], item[1], item[2], item[3], item[4], item[5],
                                                       item[6].rgb)
                file.write(line)

        print('Data saved to {}'.format(fileName))

    def savetoXML(self, fileName, baglist):
        root = ET.Element("data")
        for item in baglist:
            bag = ET.SubElement(root, "bag")
            x_first = ET.SubElement(bag, "x_first")
            x_first.text = str(item[0])
            y_first = ET.SubElement(bag, "y_first")
            y_first.text = str(item[1])
            x_second = ET.SubElement(bag, "x_second")
            x_second.text = str(item[2])
            y_second = ET.SubElement(bag, "y_second")
            y_second.text = str(item[3])
            label = ET.SubElement(bag, "label")
            label.text = item[4]
            index = ET.SubElement(bag, "index")
            index.text = str(item[5])
            color=ET.SubElement(bag,"color")
            index.text=str(item[6].rgb)

        tree = ET.ElementTree(root)
        with open(fileName, "wb") as file:
            tree.write(file)
        print('Data saved to {}'.format(fileName))

    def savetoJson(self, fileName, baglist):
        data = []
        for item in baglist:
            bag = {
                'x_first': item[0],
                'y_first': item[1],
                'x_second': item[2],
                'y_second': item[3],
                'label': item[4],
                'index': item[5],
                'color': item[6].rgb
            }
            data.append(bag)

        with open(fileName, 'w') as file:
            json.dump(data, file)

        print('Data saved to {}'.format(fileName))

    def onOpen(self):
        folderPath = QFileDialog.getExistingDirectory(self, "选择文件夹")
        if folderPath:
            self.imageList = [os.path.join(folderPath, file) for file in os.listdir(folderPath) if
                              file.lower().endswith(('.png', '.jpg', '.jpeg', '.gif'))]
            if self.imageList:
                self.showImage(self.currentIndex)
    def clearlabel(self):
        marke = markinguse()
        marke.clearData()
    def opnezgk(self):
        self.MylabeL.openzgk()

    def onOpenonly(self):
        curPath = QDir.currentPath()  # 获取系统当前目录
        title = "选择图片文件"
        filt = "图片文件(*.bmp *.png *.jpg);;所有文件(*.*)"
        fileName, flt = QFileDialog.getOpenFileName(self, title, curPath, filt)
        if (fileName == ""):
            return
        else:
            img = QPixmap(fileName)
            self.label.setPixmap(img)
            # self.label.setScaledContents(True)
            self.label.setCursor(Qt.CrossCursor)
            self.label.initParam()
            self.show()
            self.label.fileInfo = {"picturefilename": fileName,
                                   "picturebasename": os.path.basename(fileName),
                                   "picturewidth": img.width(),
                                   "pictureheight": img.height()}

    def onPrevImage(self):
        if self.currentIndex > 0:
            self.currentIndex -= 1
            self.showImage(self.currentIndex)


    def onNextImage(self):
        if self.currentIndex < len(self.imageList) - 1:
            self.currentIndex += 1
            self.showImage(self.currentIndex)

    def showImage(self, index):
        imgPath = self.imageList[index]
        img = QPixmap(imgPath)
        self.label.resetBoxes()
        self.label.setPixmap(img)
        self.label.fileInfo = {"picturefilename": imgPath,
                               "picturebasename": os.path.basename(imgPath),
                               "picturewidth": img.width(),
                               "pictureheight": img.height()}


if __name__ == '__main__':
    app = QApplication(sys.argv)
    Calc = MyMainWindow()
    qtmodern.styles.dark(app)
    mw=qtmodern.windows.ModernWindow(Calc)
    mw.setWindowTitle("图片标注小工具")
    mw.show()
    sys.exit(app.exec_())

6.总结：

啊，大家应该都能看得懂吧，代码标注都写在上面了

借鉴了不要忘了关注点赞一波哦

灌注浙财荒田鱼欧~