YOLO11深度学习 - 使用训练好的模型标注图片

-

标注图片，是一件十分麻烦事情，要真去标注一大堆图片，那一定不是合格的程序员。

先圈画个几十张，训练个模型，之后使用训练好的模型，进行图片标注。

一边升级模型，一边使用模型自动标注，人工进行校对、微调。

注意：这个过程中一定注意核对，避免混入不合格的图片标注，影响模型的质量。

from ultralytics import YOLO
import cv2
import os
from pathlib import Path

# model = YOLO("yolo11s.pt")
model = YOLO('runs/detect/coco8/weights/best.pt')


def train():
    """
    Train yolo model
    :return:
    """
    model.train(
        data='coco8.yaml',
        patience=40,  # 验证指标未改善时提前停止训练的轮数，防止过拟合
        workers=0,  # 数据加载线程数，影响数据预处理速度（多GPU训练时按RANK分配）。
        epochs=50,  # 训练轮数
        batch=16,  # 一次处理多少张，内存充足时可提升 32 或 64
        imgsz=640,  # 设置识别图片尺寸，1280 或更高
        # lr0=0.0001,      # 学习率 模型进行微调时设置
        device='cpu',
        amp=True
    )


def saveBox(image, output: str = 'output'):
    """
    识别监测区域的图片，另存为同名图片，并将内容 box 信息保存到 .txt 文件中
    :param output: 输出目录
    :param image 输入文件 "input/a.jpg"
    :return:
    """
    img = cv2.imread(image)
    results = model(img)

    filename = os.path.basename(image)

    img_height, img_width = img.shape[:2]

    # 生成对应的 txt 文件
    base_name = Path(image).stem
    txt_path = os.path.join(output, f"{base_name}.txt")

    # 处理检测结果并保存到 txt 文件
    with open(txt_path, 'w') as f:
        boxes = results[0].boxes

        # 遍历每个检测到的目标
        for box in boxes.xyxy:  # xyxy格式的边界框
            x1, y1, x2, y2 = map(int, box.tolist())  # 转换为整数坐标

            # 截取目标区域（注意 OpenCV 的坐标格式是[y:y+h, x:x+w]）
            cropped_region = img[y1:y2, x1:x2]
            output_path = os.path.join(output, f"{filename}")
            cv2.imwrite(output_path, cropped_region)  # 按坐标命名保存

            # 转换为 YOLO 格式 (x_center, y_center, width, height) 并归一化
            cls = int(boxes.cls[0])  # 类别 ID

            x_center = (x1 + x2) / 2 / img_width
            y_center = (y1 + y2) / 2 / img_height
            width = (x2 - x1) / img_width
            height = (y2 - y1) / img_height

            # 写入文件 (class x_center y_center width height)
            f.write(f"{cls} {x_center:.6f} {y_center:.6f} {width:.6f} {height:.6f}\n")


def saveAllFileBox(folder_path, output: str = 'output'):
    """
    自学习：识别图片，并保存标注数据
    :param folder_path: 图片所在的文件夹
    :param output: 输出资源的保存路径
    :return:
    """
    for root, dirs, files in os.walk(folder_path):
        for file in files:
            file_path = os.path.join(root, file)
            saveBox(file_path, output)


saveAllFileBox("D:/seaboot/yolo/input/coco8", "D:/seaboot/yolo/output/coco8")