基于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])
                    f.write(f"{label} {points}\n")
                elif isinstance(shape, QPoint):  # 点标注
                    f.write(f"{label} {shape.x()} {shape.y()}\n")
        print(f"标注已保存到: {txt_path}")

四、核心代码

class MainWindow(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("简易图像标注工具")
        self.setGeometry(100, 100, 800, 600)

        # 主布局
        self.central_widget = QWidget()
        self.setCentralWidget(self.central_widget)
        self.layout = QVBoxLayout(self.central_widget)

        # 标注类型选择
        self.annotation_type_combo = QComboBox()
        self.annotation_type_combo.addItems(["矩形框", "多边形", "点标注"])
        self.annotation_type_combo.currentTextChanged.connect(self.set_annotation_type)
        # self.layout.addWidget(self.annotation_type_combo)

        # 图片显示区域和切换按钮布局
        self.image_layout = QHBoxLayout()

        # 上一张按钮
        self.prev_button = QPushButton("上一张 (A)")
        self.prev_button.setEnabled(False)  # 初始状态禁用
        self.prev_button.setShortcut(QKeySequence("A"))  # 添加快捷键 A
        self.image_layout.addWidget(self.prev_button)

        # 图像显示区域
        self.image_label = ImageLabel(self)  # 将 MainWindow 实例传递给 ImageLabel
        self.image_layout.addWidget(self.image_label)

        # 下一张按钮
        self.next_button = QPushButton("下一张 (D)")
        self.next_button.setEnabled(False)  # 初始状态禁用
        self.next_button.setShortcut(QKeySequence("D"))  # 添加快捷键 D
        self.image_layout.addWidget(self.next_button)

        self.layout.addLayout(self.image_layout)

        # 当前图片索引和总图片数显示
        self.index_label = QLabel("0/0")
        self.index_label.setAlignment(Qt.AlignCenter)
        self.layout.addWidget(self.index_label)

        # 底部布局（按钮和标注信息显示框）
        self.bottom_layout = QHBoxLayout()

        # 标注信息显示框
        self.annotation_display = QPlainTextEdit()
        self.annotation_display.setReadOnly(True)  # 设置为只读
        self.annotation_display.setPlaceholderText("标注信息将显示在这里")
        self.bottom_layout.addWidget(self.annotation_display)

        # 按钮区域
        self.button_layout = QVBoxLayout()
        self.load_folder_button = QPushButton("加载文件夹")
        self.design_button = QPushButton("设计标签")
        self.delete_label_button = QPushButton("删除标签")
        self.save_button = QPushButton("保存标注")
        self.clear_button = QPushButton("清除标注")
        self.button_layout.addWidget(self.annotation_type_combo)
        self.button_layout.addWidget(self.load_folder_button)
        self.button_layout.addWidget(self.design_button)
        self.button_layout.addWidget(self.delete_label_button)
        self.button_layout.addWidget(self.save_button)
        self.button_layout.addWidget(self.clear_button)
        self.bottom_layout.addLayout(self.button_layout)

        self.layout.addLayout(self.bottom_layout)

        # 连接按钮事件
        self.load_folder_button.clicked.connect(self.load_folder)
        self.prev_button.clicked.connect(self.load_prev_image)
        self.next_button.clicked.connect(self.load_next_image)
        self.design_button.clicked.connect(self.design_label_types)
        self.delete_label_button.clicked.connect(self.delete_label_type)
        self.save_button.clicked.connect(self.save_annotations)
        self.clear_button.clicked.connect(self.clear_annotations)

        # 当前图片路径和文件夹信息
        self.current_image_path = None
        self.image_folder = None
        self.image_files = []
        self.current_image_index = -1

        # 标签类型
        self.label_types = []
        self.load_label_types()  # 加载已有的标签类型

    def set_annotation_type(self, annotation_type):
        """设置标注类型"""
        self.image_label.annotation_type = annotation_type

    def load_folder(self):
        """加载文件夹中的所有图片"""
        folder_path = QFileDialog.getExistingDirectory(self, "选择图片文件夹")
        if folder_path:
            self.image_folder = folder_path
            self.image_files = [
                os.path.join(folder_path, f) for f in os.listdir(folder_path)
                if f.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp'))
            ]
            if self.image_files:
                self.current_image_index = 0
                self.load_image(self.image_files[self.current_image_index])
                self.update_button_state()
                self.update_index_label()

    def load_image(self, file_path):
        """加载指定路径的图片"""
        if file_path:
            # 使用 OpenCV 读取图像
            image = cv2.imread(file_path)
            if image is not None:
                self.image_label.set_image(image)
                self.current_image_path = file_path  # 保存当前图片路径
                self.image_label.annotations = []  # 清除当前标注
                self.image_label.polygon_points = []  # 清除多边形点
                self.load_annotations()  # 加载已有的标注信息
                self.update_annotation_display()  # 更新标注信息显示

    def load_prev_image(self):
        """加载上一张图片"""
        if self.image_files and self.current_image_index > 0:
            self.current_image_index -= 1
            self.load_image(self.image_files[self.current_image_index])
            self.update_button_state()
            self.update_index_label()

    def load_next_image(self):
        """加载下一张图片"""
        if self.image_files and self.current_image_index < len(self.image_files) - 1:
            self.current_image_index += 1
            self.load_image(self.image_files[self.current_image_index])
            self.update_button_state()
            self.update_index_label()
            self.clear_annotations()

    def update_button_state(self):
        """根据当前图片索引更新按钮状态"""
        self.prev_button.setEnabled(self.current_image_index > 0)
        self.next_button.setEnabled(self.current_image_index < len(self.image_files) - 1)

    def update_index_label(self):
        """更新当前图片索引和总图片数显示"""
        self.index_label.setText(f"{self.current_image_index + 1}/{len(self.image_files)}")

    def design_label_types(self):
        """设计标签类型"""
        label, ok = QInputDialog.getText(self, "添加标签类型", "请输入新的标签类型:")
        if ok and label:
            if label not in self.label_types:
                self.label_types.append(label)
                self.save_label_types()  # 保存标签类型
                QMessageBox.information(self, "成功", f"标签 '{label}' 已添加！")
            else:
                QMessageBox.warning(self, "警告", f"标签 '{label}' 已存在！")

    def delete_label_type(self):
        """删除标签类型"""
        if self.label_types:
            label, ok = QInputDialog.getItem(self, "删除标签类型", "请选择要删除的标签类型:", self.label_types, 0, False)
            if ok and label:
                self.label_types.remove(label)
                self.save_label_types()  # 保存标签类型
                QMessageBox.information(self, "成功", f"标签 '{label}' 已删除！")
        else:
            QMessageBox.warning(self, "警告", "没有可删除的标签类型！")

    def save_label_types(self):
        """保存标签类型到文件"""
        with open("label_types.txt", "w") as f:
            for label in self.label_types:
                f.write(f"{label}\n")

    def load_label_types(self):
        """从文件加载标签类型"""
        if os.path.exists("label_types.txt"):
            with open("label_types.txt", "r") as f:
                self.label_types = [line.strip() for line in f.readlines()]

    def load_annotations(self):
        """加载已有的标注信息"""
        if self.current_image_path is None:
            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")

        # 清空当前标注
        self.image_label.annotations = []

        # 加载标注信息
        if os.path.exists(txt_path):
            with open(txt_path, "r") as f:
                for line in f:
                    parts = line.strip().split()
                    if len(parts) == 5:
                        label, xmin, ymin, xmax, ymax = parts
                        rect = QRect(int(xmin), int(ymin), int(xmax) - int(xmin), int(ymax) - int(ymin))
                        self.image_label.annotations.append((rect, label))
            self.image_label.update_display()

    def save_annotations(self):
        """保存标注结果"""
        if not self.image_label.annotations:
            QMessageBox.information(self, "信息", "没有标注可保存")
            return

        if self.current_image_path is None:
            QMessageBox.warning(self, "警告", "未加载图片")
            return

        try:
            # 默认保存路径为图片的同路径
            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])
                        f.write(f"{label} {points}\n")
                    elif isinstance(shape, QPoint):  # 点标注
                        f.write(f"{label} {shape.x()} {shape.y()}\n")
            QMessageBox.information(self, "信息", f"标注已保存到: {txt_path}")
        except Exception as e:
            QMessageBox.critical(self, "错误", f"保存标注时出错: {str(e)}")

    def clear_annotations(self):
        """清除所有标注"""
        self.image_label.annotations = []
        self.image_label.update_display()
        self.update_annotation_display()  # 更新标注信息显示

    def update_annotation_display(self):
        """更新标注信息显示框"""
        self.annotation_display.clear()
        for shape, label in self.image_label.annotations:
            if isinstance(shape, QRect):  # 矩形框
                self.annotation_display.appendPlainText(
                    f"{label}: 矩形框 ({shape.left()}, {shape.top()}, {shape.right()}, {shape.bottom()})"
                )
            elif isinstance(shape, QPolygon):  # 多边形
                points = ", ".join([f"({point.x()}, {point.y()})" for point in shape])
                self.annotation_display.appendPlainText(f"{label}: 多边形 [{points}]")
            elif isinstance(shape, QPoint):  # 点标注
                self.annotation_display.appendPlainText(f"{label}: 点 ({shape.x()}, {shape.y()})")

五、源码地址

LabelTool
未来将基于YOLOv8训练一个识别图像内瑕疵的模型。