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

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()

运行截图：