YoloV8图片标注工具

import sys
import os
from PyQt6.QtWidgets import (
    QApplication, QMainWindow, QWidget, QVBoxLayout, QHBoxLayout,
    QPushButton, QFileDialog, QListWidget, QDialog,
    QDialogButtonBox, QLabel, QComboBox, QMessageBox, QLineEdit
)
from PyQt6.QtGui import QPixmap, QPainter, QPen, QColor
from PyQt6.QtCore import Qt, QPoint, QRect


class ClassSelectDialog(QDialog):
    def __init__(self, class_list, parent=None):
        super().__init__(parent)
        self.setWindowTitle("选择目标类别")
        self.resize(300, 120)
        layout = QVBoxLayout()
        self.combo = QComboBox()
        self.combo.addItems(class_list)
        layout.addWidget(QLabel("请选择类别："))
        layout.addWidget(self.combo)
        btn_box = QDialogButtonBox(
            QDialogButtonBox.StandardButton.Ok | QDialogButtonBox.StandardButton.Cancel
        )
        btn_box.accepted.connect(self.accept)
        btn_box.rejected.connect(self.reject)
        layout.addWidget(btn_box)
        self.setLayout(layout)

    def get_selected_class(self):
        return self.combo.currentText()


class ImageCanvas(QLabel):
    def __init__(self, parent):
        super().__init__(parent)
        self.parent_win = parent
        self.setAlignment(Qt.AlignmentFlag.AlignCenter)
        self.setAutoFillBackground(True)

        self.img_pix = None
        self.scaled_pix = None
        self.start_pt = None
        self.drawing_rect = None
        self.box_list = []

    def load_image(self, img_path):
        self.img_pix = QPixmap(img_path)
        self.box_list.clear()
        self.update_scaled_pixmap()
        self.update()

    def resizeEvent(self, event):
        super().resizeEvent(event)
        self.update_scaled_pixmap()

    def update_scaled_pixmap(self):
        if self.img_pix is None or self.img_pix.isNull():
            self.scaled_pix = None
            self.setPixmap(QPixmap())
            return

        self.scaled_pix = self.img_pix.scaled(
            self.size(),
            Qt.AspectRatioMode.KeepAspectRatio,
            Qt.TransformationMode.SmoothTransformation
        )
        self.setPixmap(self.scaled_pix)

    def display_rect(self):
        if self.scaled_pix is None or self.scaled_pix.isNull():
            return QRect(0, 0, 0, 0)

        label_w = self.width()
        label_h = self.height()
        pix_w = self.scaled_pix.width()
        pix_h = self.scaled_pix.height()
        x = (label_w - pix_w) // 2
        y = (label_h - pix_h) // 2
        return QRect(x, y, pix_w, pix_h)

    def mouse_to_image(self, pos: QPoint):
        drect = self.display_rect()
        if drect.width() == 0 or drect.height() == 0:
            return None
        if not drect.contains(pos):
            return None

        sx = (pos.x() - drect.x()) / drect.width()
        sy = (pos.y() - drect.y()) / drect.height()
        orig_x = sx * self.img_pix.width()
        orig_y = sy * self.img_pix.height()
        return orig_x, orig_y

    def image_to_display(self, orig_x, orig_y):
        drect = self.display_rect()
        if drect.width() == 0 or drect.height() == 0:
            return None

        sx = orig_x / self.img_pix.width()
        sy = orig_y / self.img_pix.height()
        x = drect.x() + sx * drect.width()
        y = drect.y() + sy * drect.height()
        return x, y

    def paintEvent(self, event):
        painter = QPainter(self)

        # 白色背景
        painter.fillRect(self.rect(), QColor(255, 255, 255))

        # 绘制原图
        if self.scaled_pix and not self.scaled_pix.isNull():
            drect = self.display_rect()
            painter.drawPixmap(drect, self.scaled_pix)

        # 已确认标注框：红色实线
        if self.img_pix and self.scaled_pix:
            pen_saved = QPen(QColor(220, 30, 30), 2)
            painter.setPen(pen_saved)

            for x1, y1, x2, y2, cls_name in self.box_list:
                p1 = self.image_to_display(x1, y1)
                p2 = self.image_to_display(x2, y2)
                if p1 is None or p2 is None:
                    continue

                rect = QRect(
                    QPoint(int(p1[0]), int(p1[1])),
                    QPoint(int(p2[0]), int(p2[1]))
                )
                painter.drawRect(rect)
                painter.drawText(rect.topLeft() + QPoint(0, -4), cls_name)

        # 正在拖拽的临时框：绿色虚线
        if self.start_pt and self.drawing_rect:
            pen_temp = QPen(QColor(0, 180, 0), 2, Qt.PenStyle.DashLine)
            painter.setPen(pen_temp)
            painter.drawRect(self.drawing_rect)

        painter.end()

    def mousePressEvent(self, event):
        if self.img_pix is None:
            return
        if event.button() == Qt.MouseButton.LeftButton:
            self.start_pt = event.position().toPoint()
            self.drawing_rect = QRect(self.start_pt, self.start_pt)

    def mouseMoveEvent(self, event):
        if self.start_pt is not None:
            end_pt = event.position().toPoint()
            self.drawing_rect = QRect(self.start_pt, end_pt).normalized()
            self.update()

    def mouseReleaseEvent(self, event):
        if event.button() == Qt.MouseButton.LeftButton and self.start_pt is not None:
            end_pt = event.position().toPoint()
            temp_rect = QRect(self.start_pt, end_pt).normalized()
            self.start_pt = None
            self.drawing_rect = None

            if temp_rect.width() > 8 and temp_rect.height() > 8:
                p1 = self.mouse_to_image(temp_rect.topLeft())
                p2 = self.mouse_to_image(temp_rect.bottomRight())
                if p1 is not None and p2 is not None:
                    x1, y1 = p1
                    x2, y2 = p2
                    dlg = ClassSelectDialog(self.parent_win.class_names, self)
                    if dlg.exec():
                        cls = dlg.get_selected_class()
                        self.box_list.append((x1, y1, x2, y2, cls))

            self.update()


class MainAnnotateWindow(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("YOLO图片标注工具 v2.2")
        self.resize(1050, 720)
        self.class_names = []
        self.current_img_path = ""
        # 项目根目录：脚本所在文件夹
        self.root_dir = os.path.dirname(os.path.abspath(__file__))

        central_widget = QWidget()
        self.setCentralWidget(central_widget)
        main_layout = QHBoxLayout(central_widget)

        left_panel = QWidget()
        left_panel.setFixedWidth(230)
        left_layout = QVBoxLayout(left_panel)

        self.btn_import = QPushButton("导入图片")
        self.btn_add_cls = QPushButton("添加目标分类")
        self.list_cls = QListWidget()
        self.btn_save_txt = QPushButton("保存批注(txt)")

        left_layout.addWidget(self.btn_import)
        left_layout.addWidget(self.btn_add_cls)
        left_layout.addWidget(QLabel("类别列表："))
        left_layout.addWidget(self.list_cls)
        left_layout.addWidget(self.btn_save_txt)
        left_layout.addStretch()

        self.canvas = ImageCanvas(self)
        main_layout.addWidget(left_panel)
        main_layout.addWidget(self.canvas)

        self.btn_import.clicked.connect(self.on_import_img)
        self.btn_add_cls.clicked.connect(self.on_add_class)
        self.btn_save_txt.clicked.connect(self.on_save_annotation)

    def on_add_class(self):
        dlg = QDialog(self)
        dlg.setWindowTitle("新增分类")
        lay = QVBoxLayout(dlg)
        edit = QLineEdit()
        lay.addWidget(QLabel("输入类别名称："))
        lay.addWidget(edit)
        btns = QDialogButtonBox(
            QDialogButtonBox.StandardButton.Ok | QDialogButtonBox.StandardButton.Cancel
        )
        btns.accepted.connect(dlg.accept)
        btns.rejected.connect(dlg.reject)
        lay.addWidget(btns)
        if dlg.exec():
            name = edit.text().strip()
            if name and name not in self.class_names:
                self.class_names.append(name)
                self.list_cls.addItem(name)

    def on_import_img(self):
        path, _ = QFileDialog.getOpenFileName(
            self, "选择图片", "", "图像文件(*.jpg *.jpeg *.png *.bmp)"
        )
        if not path:
            return
        self.current_img_path = path
        self.canvas.load_image(path)

    def on_save_annotation(self):
        if not self.current_img_path:
            QMessageBox.warning(self, "提示", "请先导入图片！")
            return
        if not self.canvas.box_list:
            QMessageBox.warning(self, "提示", "当前没有标注框，无需保存！")
            return
        if not self.class_names:
            QMessageBox.warning(self, "提示", "请先添加目标分类！")
            return

        from PyQt6.QtGui import QImage
        img = QImage(self.current_img_path)
        img_w = img.width()
        img_h = img.height()
        base_filename = os.path.splitext(os.path.basename(self.current_img_path))[0]
        # 修改：保存到脚本所在项目根目录
        save_txt_path = os.path.join(self.root_dir, f"{base_filename}.txt")

        output_lines = []
        for x1, y1, x2, y2, cls_name in self.canvas.box_list:
            cls_id = self.class_names.index(cls_name)
            box_w = x2 - x1
            box_h = y2 - y1
            x_center = (x1 + box_w / 2) / img_w
            y_center = (y1 + box_h / 2) / img_h
            w_norm = box_w / img_w
            h_norm = box_h / img_h
            line = f"{cls_id} {x_center:.6f} {y_center:.6f} {w_norm:.6f} {h_norm:.6f}"
            output_lines.append(line)

        with open(save_txt_path, "w", encoding="utf-8") as f:
            f.write("\n".join(output_lines))
        QMessageBox.information(self, "保存成功", f"标注文件已保存到项目根目录：\n{save_txt_path}")


if __name__ == "__main__":
    app = QApplication(sys.argv)
    win = MainAnnotateWindow()
    win.show()
    sys.exit(app.exec())