模仿精灵助手，使用pyqt5完成数据标注功能

import sys
import os
from PyQt5.QtWidgets import (QApplication, QMainWindow, QWidget, QVBoxLayout,
                             QHBoxLayout, QPushButton, QLabel, QFileDialog,
                             QListWidget, QInputDialog, QMessageBox)
from PyQt5.QtGui import QPixmap, QPainter, QPen, QColor, QBrush, QFont
from PyQt5.QtCore import Qt, QPoint, QRect
from xml.etree import ElementTree as ET
from xml.dom import minidom

class AnnotationTool(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("目标检测标注工具")
        self.setFixedSize(1200, 800)
        self.setWindowFlags(self.windowFlags() & ~Qt.WindowMaximizeButtonHint)
        
        self.image_path = ""
        self.original_pixmap = None
        self.start_point = QPoint()
        self.end_point = QPoint()
        self.drawing = False
        self.annotations = []
        self.current_scaled_pixmap = None
        
        central_widget = QWidget()
        self.setCentralWidget(central_widget)
        
        main_layout = QHBoxLayout(central_widget)
        
        left_panel = QWidget()
        left_layout = QVBoxLayout(left_panel)
        left_panel.setFixedWidth(200)
        
        self.open_btn = QPushButton("打开图片")
        self.save_btn = QPushButton("保存标注")
        self.add_class_btn = QPushButton("添加类别")
        self.clear_btn = QPushButton("清空标注")
        
        self.open_btn.clicked.connect(self.open_image)
        self.save_btn.clicked.connect(self.save_annotation)
        self.add_class_btn.clicked.connect(self.add_class)
        self.clear_btn.clicked.connect(self.clear_annotations)
        
        self.class_list = QListWidget()
        default_classes = ["Person", "Car", "Bicycle", "Dog", "Cat"]
        self.class_list.addItems(default_classes)
        
        self.annotation_list = QListWidget()
        
        left_layout.addWidget(QLabel("物体类别"))
        left_layout.addWidget(self.class_list)
        left_layout.addWidget(self.add_class_btn)
        left_layout.addSpacing(20)
        left_layout.addWidget(QLabel("标注列表"))
        left_layout.addWidget(self.annotation_list)
        left_layout.addSpacing(20)
        left_layout.addWidget(self.open_btn)
        left_layout.addWidget(self.save_btn)
        left_layout.addWidget(self.clear_btn)
        left_layout.addStretch()
        
        self.image_label = QLabel()
        self.image_label.setAlignment(Qt.AlignCenter)
        self.image_label.setStyleSheet("border: 1px solid gray;")
        self.image_label.setFixedSize(980, 780)
        
        main_layout.addWidget(left_panel)
        main_layout.addWidget(self.image_label)
        
    def open_image(self):
        file_path, _ = QFileDialog.getOpenFileName(
            self, "选择图片", "",
            "Image Files (*.png *.jpg *.jpeg *.bmp *.tif)"
        )
        if file_path:
            self.image_path = file_path
            self.original_pixmap = QPixmap(file_path)
            self.current_scaled_pixmap = self.original_pixmap.scaled(
                self.image_label.size(),
                Qt.KeepAspectRatio,
                Qt.SmoothTransformation
            )
            self.image_label.setPixmap(self.current_scaled_pixmap)
            self.annotations.clear()
            self.annotation_list.clear()
            
    def mousePressEvent(self, event):
        if event.button() == Qt.LeftButton and self.original_pixmap:
            label_pos = self.image_label.mapFromGlobal(event.globalPos())
            if self.image_label.rect().contains(label_pos):
                self.start_point = self.get_image_coordinate(event.pos())
                self.drawing = True
            
    def mouseMoveEvent(self, event):
        if self.drawing and self.original_pixmap:
            self.end_point = self.get_image_coordinate(event.pos())
            self.draw_rectangle()
            
    def mouseReleaseEvent(self, event):
        if event.button() == Qt.LeftButton and self.drawing and self.original_pixmap:
            self.drawing = False
            self.end_point = self.get_image_coordinate(event.pos())
            
            xmin = min(self.start_point.x(), self.end_point.x())
            ymin = min(self.start_point.y(), self.end_point.y())
            xmax = max(self.start_point.x(), self.end_point.x())
            ymax = max(self.start_point.y(), self.end_point.y())
            
            if (xmax - xmin) > 5 and (ymax - ymin) > 5:
                current_item = self.class_list.currentItem()
                label = current_item.text() if current_item else (self.class_list.item(0).text() if self.class_list.count() > 0 else "Unknown")
                
                self.annotations.append((xmin, ymin, xmax, ymax, label))
                self.annotation_list.addItem(f"{label}: ({xmin}, {ymin}) - ({xmax}, {ymax})")
                self.draw_all_annotations()
                
    def get_image_coordinate(self, pos):
        if not self.original_pixmap or not self.current_scaled_pixmap:
            return QPoint()
            
        label_rect = self.image_label.rect()
        scaled_rect = self.current_scaled_pixmap.rect()
        
        offset_x = (label_rect.width() - scaled_rect.width()) // 2
        offset_y = (label_rect.height() - scaled_rect.height()) // 2
        
        scaled_x = pos.x() - self.image_label.pos().x() - offset_x
        scaled_y = pos.y() - self.image_label.pos().y() - offset_y
        
        scaled_x = max(0, min(scaled_x, scaled_rect.width()))
        scaled_y = max(0, min(scaled_y, scaled_rect.height()))
        
        scale_x = self.original_pixmap.width() / scaled_rect.width()
        scale_y = self.original_pixmap.height() / scaled_rect.height()
        
        real_x = int(scaled_x * scale_x)
        real_y = int(scaled_y * scale_y)
        
        real_x = max(0, min(real_x, self.original_pixmap.width()))
        real_y = max(0, min(real_y, self.original_pixmap.height()))
        
        return QPoint(real_x, real_y)
        
    def draw_rectangle(self):
        temp_pixmap = self.original_pixmap.copy()
        painter = QPainter(temp_pixmap)
        painter.setRenderHint(QPainter.Antialiasing)
        
        self.draw_existing_annotations(painter)
        
        painter.setPen(QPen(QColor(255, 0, 0), 2, Qt.SolidLine))
        painter.setBrush(Qt.NoBrush)
        rect = QRect(self.start_point, self.end_point).normalized()
        painter.drawRect(rect)
        
        painter.end()
        
        self.current_scaled_pixmap = temp_pixmap.scaled(
            self.image_label.size(),
            Qt.KeepAspectRatio,
            Qt.SmoothTransformation
        )
        self.image_label.setPixmap(self.current_scaled_pixmap)
        
    def draw_all_annotations(self):
        if not self.original_pixmap:
            return
            
        temp_pixmap = self.original_pixmap.copy()
        painter = QPainter(temp_pixmap)
        painter.setRenderHint(QPainter.Antialiasing)
        
        self.draw_existing_annotations(painter)
        
        painter.end()
        
        self.current_scaled_pixmap = temp_pixmap.scaled(
            self.image_label.size(),
            Qt.KeepAspectRatio,
            Qt.SmoothTransformation
        )
        self.image_label.setPixmap(self.current_scaled_pixmap)
        
    def draw_existing_annotations(self, painter):
        for (xmin, ymin, xmax, ymax, label) in self.annotations:
            color = QColor.fromHsv(((hash(label) % 360)), 200, 255)
            
            painter.setPen(QPen(color, 2, Qt.SolidLine))
            painter.setBrush(Qt.NoBrush)
            painter.drawRect(xmin, ymin, xmax - xmin, ymax - ymin)
            
            font = QFont()
            font.setBold(True)
            painter.setFont(font)
            
            text_rect = painter.boundingRect(QRect(0, 0, 200, 30), Qt.AlignLeft, label)
            text_width = text_rect.width() + 8
            text_height = text_rect.height() + 4
            
            text_bg_x = xmin
            text_bg_y = max(0, ymin - text_height)
            
            painter.save()
            painter.setPen(Qt.NoPen)
            painter.setBrush(QBrush(color))
            painter.drawRect(text_bg_x, text_bg_y, text_width, text_height)
            painter.restore()
            
            painter.setPen(QPen(Qt.white))
            text_y = text_bg_y + text_height - 4
            painter.drawText(text_bg_x + 4, text_y, label)
        
    def add_class(self):
        class_name, ok = QInputDialog.getText(self, "添加类别", "输入新类别名称:")
        if ok and class_name.strip():
            exists = any(self.class_list.item(i).text() == class_name for i in range(self.class_list.count()))
            if not exists:
                self.class_list.addItem(class_name)
            else:
                QMessageBox.warning(self, "警告", "该类别已存在！")
                
    def clear_annotations(self):
        self.annotations.clear()
        self.annotation_list.clear()
        if self.original_pixmap:
            self.current_scaled_pixmap = self.original_pixmap.scaled(
                self.image_label.size(),
                Qt.KeepAspectRatio,
                Qt.SmoothTransformation
            )
            self.image_label.setPixmap(self.current_scaled_pixmap)
            
    def save_annotation(self):
        if not self.image_path or not self.annotations:
            QMessageBox.warning(self, "警告", "没有可保存的标注数据！")
            return
            
        save_path, _ = QFileDialog.getSaveFileName(
            self, "保存标注文件",
            os.path.splitext(self.image_path)[0] + ".xml",
            "XML Files (*.xml)"
        )
        if save_path:
            self.create_xml_annotation(save_path)
            QMessageBox.information(self, "成功", "标注文件已保存！")
            
    def create_xml_annotation(self, save_path):
        annotation = ET.Element("annotation")
        
        folder = ET.SubElement(annotation, "folder")
        folder.text = os.path.dirname(self.image_path)
        
        filename = ET.SubElement(annotation, "filename")
        filename.text = os.path.basename(self.image_path)
        
        path = ET.SubElement(annotation, "path")
        path.text = self.image_path
        
        size = ET.SubElement(annotation, "size")
        width = ET.SubElement(size, "width")
        width.text = str(self.original_pixmap.width())
        height = ET.SubElement(size, "height")
        height.text = str(self.original_pixmap.height())
        depth = ET.SubElement(size, "depth")
        depth.text = "3"
        
        for (xmin, ymin, xmax, ymax, label) in self.annotations:
            obj = ET.SubElement(annotation, "object")
            name = ET.SubElement(obj, "name")
            name.text = label
            bndbox = ET.SubElement(obj, "bndbox")
            ET.SubElement(bndbox, "xmin").text = str(xmin)
            ET.SubElement(bndbox, "ymin").text = str(ymin)
            ET.SubElement(bndbox, "xmax").text = str(xmax)
            ET.SubElement(bndbox, "ymax").text = str(ymax)
        
        xml_str = minidom.parseString(ET.tostring(annotation)).toprettyxml(indent="  ")
        xml_str = '\n'.join([line for line in xml_str.split('\n') if line.strip()])
        with open(save_path, "w", encoding="utf-8") as f:
            f.write(xml_str)

if __name__ == "__main__":
    app = QApplication(sys.argv)
    window = AnnotationTool()
    window.show()
    sys.exit(app.exec_())