基于YOLOv8模型训练之图像特征标注程序

一、引言

在着手训练自定义模型的过程中，我们不可避免地需要准备大量的图像（image）及其对应的标签（label）信息。这些标签的精准程度，对于模型的精确率（Precision）、准确率（Accuracy）以及F1值等关键性能指标具有直接且深远的影响。鉴于此，本人专门设计并开发了一款图像特征标注程序，旨在高效辅助模型的训练工作。
程序无偿分享，欢迎修改指正！

二、UI 界面

三、功能简介

该程序的主要功能包括：通过选择图片文件夹来预览图像；在设置标注类型后（再次启动程序后无需重新设置，对图片中的特征进行标注并记录；完成每张图片的标注后，点击保存并继续下一张图片的标注工作；同时，程序还支持对标签类型进行修改。这些功能共同构成了该图像特征标注程序的核心用途。

1. 核心类库

PySide6：负责程序的界面的搭建；
python-opencv：负责每张图片的标注；

2. 图像特征标注

首先，我创建了一个名为ImageLabel的类，专门负责处理每张图片的显示（800x600 图片尺寸较大将会被缩放到规定尺寸）及其特征标注功能。该类默认采用矩形框作为标注特征的方式，并通过监听鼠标事件来精确捕捉并标注图片内的特征点。

        # 将 OpenCV 图像转换为 QImage
        height, width, channel = self.scaled_image.shape
        bytes_per_line = 3 * width
        q_img = QImage(self.scaled_image.data, width, height, bytes_per_line, QImage.Format_RGB888).rgbSwapped()

        # 绘制标注
        painter = QPainter(q_img)
        pen = QPen(QColor(255, 0, 0), 2)  # 红色画笔
        painter.setPen(pen)
        for annotation in self.annotations:
            shape, label = annotation
            if isinstance(shape, QRect):  # 矩形框
                painter.drawRect(shape)
                painter.drawText(shape.topLeft(), label)
            elif isinstance(shape, QPolygon):  # 多边形
                painter.drawPolygon(shape)
                painter.drawText(shape.boundingRect().topLeft(), label)
            elif isinstance(shape, QPoint):  # 点标注
                painter.drawPoint(shape)
                painter.drawText(shape, label)
        if self.current_shape:
            if isinstance(self.current_shape, QRect):  # 当前绘制的矩形框
                painter.drawRect(self.current_shape)
            elif isinstance(self.current_shape, QPolygon):  # 当前绘制的多边形
                painter.drawPolygon(self.current_shape)
            elif isinstance(self.current_shape, QPoint):  # 当前绘制的点
                painter.drawPoint(self.current_shape)
        painter.end()

        # 显示图像
        self.setPixmap(QPixmap.fromImage(q_img).scaled(self.max_size, Qt.KeepAspectRatio, Qt.SmoothTransformation))

3. 图像特征存储

通过捕获鼠标的位置坐标，确定矩形框的左上顶点坐标和右下端点坐标。随后，创建一个与图片具有相同名称的 .txt 文件，并将所选的标记类型（例如“stone”）以及对应的标记框信息写入该文件中，每一行记录一个标记框的信息。

    def save_annotations(self):
        """保存标注结果"""
        if not self.image_label.annotations:
            print("没有标注可保存")
            return

        if self.current_image_path is None:
            print("未加载图片")
            return

        # 默认保存路径为图片的同路径
        image_name = os.path.splitext(os.path.basename(self.current_image_path))[0]
        txt_path = os.path.join(self.image_folder, f"{
     image_name}_annotations.txt")

        # 保存标注信息到 TXT 文件
        with open(txt_path, "w") as f:
            for shape, label in self.image_label.annotations:
                if isinstance(shape, QRect):  # 矩形框
                    f.write(f"{
     label} {
     shape.left()} {
     shape.top()} {
     shape.right()} {
     shape.bottom()}\n")
                elif isinstance(shape, QPolygon):  # 多边形
                    points = ", ".join([f"{
     point.x()} {
     point.y()}" for point in shape]