自制图形图像注释工具-用于生成xml文件

2024-01-07 18:25:18
import sys
import xml.etree.ElementTree as ET
from PyQt5.QtWidgets import QApplication, QMainWindow, QPushButton, QLabel, QVBoxLayout, QWidget, QFileDialog, QLineEdit, QHBoxLayout, QMessageBox
from PyQt5.QtGui import QPixmap, QPainter, QPen
from PyQt5.QtCore import Qt, QPoint, QRect

class ImageLabel(QLabel):
    def __init__(self):
        super().__init__()
        self.image = None
        self.start_point = QPoint()
        self.end_point = QPoint()
        self.drawing = False
        self.rectangles = []  # List to store rectangles
        self.current_label = None
        self.undo_stack = []  # Stack for undo operations
        self.redo_stack = []  # Stack for redo operations

    def set_image(self, image_path):
        self.image = QPixmap(image_path)
        self.setPixmap(self.image)
        self.rectangles.clear()  # Clear previous rectangles
        self.undo_stack.clear()  # Clear undo stack
        self.redo_stack.clear()  # Clear redo stack

    def set_current_label(self, label):
        self.current_label = label

    def paintEvent(self, event):
        super().paintEvent(event)
        if self.image:
            painter = QPainter(self)
            painter.drawPixmap(self.rect(), self.image)
            pen = QPen(Qt.red, 2, Qt.SolidLine)
            painter.setPen(pen)
            for rect in self.rectangles:
                painter.drawRect(rect["rectangle"])

            if self.drawing:
                painter.drawRect(QRect(self.start_point, self.end_point))

    def mousePressEvent(self, event):
        self.drawing = True
        self.start_point = event.pos()
        self.end_point = event.pos()
        self.update()

    def mouseMoveEvent(self, event):
        if self.drawing:
            self.end_point = event.pos()
            self.update()

    def mouseReleaseEvent(self, event):
        if self.current_label and self.drawing:
            rect = QRect(self.start_point, self.end_point)
            rectangle = {"label": self.current_label, "rectangle": rect}
            self.rectangles.append(rectangle)
            self.undo_stack.append(("Add", rectangle))
            self.redo_stack.clear()
        self.drawing = False
        self.update()

    def undo(self):
        if self.undo_stack:
            action, rectangle = self.undo_stack.pop()
            if action == "Add":
                self.rectangles.remove(rectangle)
                self.redo_stack.append(("Remove", rectangle))
            self.update()

    def redo(self):
        if self.redo_stack:
            action, rectangle = self.redo_stack.pop()
            if action == "Remove":
                self.rectangles.append(rectangle)
                self.undo_stack.append(("Add", rectangle))
            self.update()

class LabelImgApp(QMainWindow):
    def __init__(self):
        super().__init__()
        self.initUI()

    def initUI(self):
        self.setWindowTitle('LabelImg - Extended')
        self.setGeometry(100, 100, 800, 600)

        self.image_label = ImageLabel()

        load_button = QPushButton('Load Image', self)
        load_button.clicked.connect(self.loadImage)

        self.label_input = QLineEdit(self)
        self.label_input.setPlaceholderText("Enter label here")

        set_label_button = QPushButton('Set Label', self)
        set_label_button.clicked.connect(self.setLabel)

        save_button = QPushButton('Save Annotations', self)
        save_button.clicked.connect(self.saveAnnotations)

        undo_button = QPushButton('Undo', self)
        undo_button.clicked.connect(self.undo)

        redo_button = QPushButton('Redo', self)
        redo_button.clicked.connect(self.redo)

        layout = QVBoxLayout()
        layout.addWidget(self.image_label)

        label_layout = QHBoxLayout()
        label_layout.addWidget(self.label_input)
        label_layout.addWidget(set_label_button)

        button_layout = QHBoxLayout()
        button_layout.addWidget(load_button)
        button_layout.addWidget(save_button)
        button_layout.addWidget(undo_button)
        button_layout.addWidget(redo_button)

        layout.addLayout(label_layout)
        layout.addLayout(button_layout)

        container = QWidget()
        container.setLayout(layout)
        self.setCentralWidget(container)

    def loadImage(self):
        options = QFileDialog.Options()
        supportedFormats = "JPEG (*.jpg *.jpeg);;PNG (*.png);;BMP (*.bmp);;GIF (*.gif);;TIFF (*.tiff *.tif);;All Files (*)"
        file_name, _ = QFileDialog.getOpenFileName(self, "Open Image", "", supportedFormats, options=options)
        if file_name:
            self.image_label.set_image(file_name)

    def setLabel(self):
        label = self.label_input.text()
        if label:
            self.image_label.set_current_label(label)
        else:
            QMessageBox.warning(self, "No Label", "Please enter a label.")

    def saveAnnotations(self):
        annotations = [{"label": rect["label"],
                        "rectangle": [rect["rectangle"].x(), rect["rectangle"].y(),
                                      rect["rectangle"].width(), rect["rectangle"].height()]}
                       for rect in self.image_label.rectangles]

        if annotations:
            options = QFileDialog.Options()
            file_name, _ = QFileDialog.getSaveFileName(self, "Save File", "",
                                                       "XML Files (*.xml);;All Files (*)",
                                                       options=options)
            if file_name:
                if not file_name.endswith('.xml'):
                    file_name += '.xml'
                tree = self.annotations_to_xml(annotations)
                tree.write(file_name)
                QMessageBox.information(self, "Saved", "Annotations saved successfully.")
        else:
            QMessageBox.warning(self, "No Annotations", "There are no annotations to save.")

    def annotations_to_xml(self, annotations):
        root = ET.Element("annotations")
        for ann in annotations:
            ann_element = ET.SubElement(root, "annotation")
            ET.SubElement(ann_element, "label").text = ann["label"]
            ET.SubElement(ann_element, "x").text = str(ann["rectangle"][0])
            ET.SubElement(ann_element, "y").text = str(ann["rectangle"][1])
            ET.SubElement(ann_element, "width").text = str(ann["rectangle"][2])
            ET.SubElement(ann_element, "height").text = str(ann["rectangle"][3])
        tree = ET.ElementTree(root)
        return tree

    def undo(self):
        self.image_label.undo()

    def redo(self):
        self.image_label.redo()

def main():
    app = QApplication(sys.argv)
    ex = LabelImgApp()
    ex.show()
    sys.exit(app.exec_())

if __name__ == '__main__':
    main()

运行截图:

文章来源:https://blog.csdn.net/qq_60245590/article/details/135432813
本文来自互联网用户投稿,该文观点仅代表作者本人,不代表本站立场。本站仅提供信息存储空间服务,不拥有所有权,不承担相关法律责任。