求推荐基于PyTorch的YOLOv7图像数据集自动化标注工具或脚本

1. 批量自动标注Python脚本 #

直接利用YOLOv7预训练模型批量处理图像，生成YOLO格式的标注文件（每个图像对应同名txt，每行格式为类别ID x_center y_center width height）。

import torch
import os
from PIL import Image
from pathlib import Path

# 加载YOLOv7预训练模型（可替换为自定义微调后的权重）
model = torch.hub.load('WongKinYiu/yolov7', 'yolov7', pretrained=True)
model.eval()
device = 'cuda' if torch.cuda.is_available() else 'cpu'
model.to(device)

# 配置路径与参数
IMAGE_DIR = "./dataset/images"  # 待标注图像目录
LABEL_DIR = "./dataset/labels"  # 标注文件输出目录
CONF_THRESHOLD = 0.5  # 置信度阈值，过滤低置信度结果
CLASS_NAMES = model.names  # 模型自带类别列表

# 创建标注目录
os.makedirs(LABEL_DIR, exist_ok=True)

# 遍历处理所有图像
for img_path in Path(IMAGE_DIR).glob("*.jpg"):  # 可修改为png等其他格式
    img = Image.open(img_path).convert("RGB")
    img_w, img_h = img.size

    # 模型推理
    with torch.no_grad():
        results = model(img)

    # 生成YOLO格式标注
    label_file = os.path.join(LABEL_DIR, f"{img_path.stem}.txt")
    with open(label_file, "w") as f:
        for det in results.xyxy[0]:  # xyxy格式：x1, y1, x2, y2, conf, class_id
            conf = det[4].item()
            cls_id = int(det[5].item())
            if conf < CONF_THRESHOLD:
                continue
            # 转换为归一化坐标
            x_center = (det[0] + det[2]) / 2 / img_w
            y_center = (det[1] + det[3]) / 2 / img_h
            width = (det[2] - det[0]) / img_w
            height = (det[3] - det[1]) / img_h
            # 写入标注
            f.write(f"{cls_id} {x_center:.6f} {y_center:.6f} {width:.6f} {height:.6f}\n")
    print(f"标注完成：{label_file}")

脚本说明 #

• 依赖安装：执行pip install torch torchvision pillow即可，torch.hub会自动拉取YOLOv7代码
• 权重替换：如果有自定义微调的YOLOv7权重，可将pretrained=True改为weights='./path/to/your/model.pt'
• 阈值调整：根据数据集场景修改CONF_THRESHOLD，阈值越高标注精度越高但可能漏检，反之则召回率更高

2. 半自动标注工具 #

若需要对自动标注结果进行修正（避免误标、漏标），可使用以下工具：

• LabelStudio：集成YOLOv7作为预标注器，先自动生成标注框，再手动调整，支持直接导出YOLO格式标注
• YOLOv7官方detect.py：修改官方检测脚本，添加标注文件保存逻辑，官方脚本原生支持批量检测，只需将检测结果转换为YOLO格式写入文件即可
• 增强版LabelImg：第三方修改的LabelImg集成了YOLO预检测功能，标注时自动框出目标，可直接在界面上调整或删除框

3. 精度优化建议 #

• 小批量测试：先在少量图像上测试自动标注结果，调整置信度阈值至合适范围后再批量处理
• 模型微调：如果是特定领域数据集（如工业缺陷、医疗影像），先用少量手动标注数据微调YOLOv7，再用微调后的模型做自动标注，准确率会大幅提升
• 后处理优化：添加NMS（非极大值抑制）过滤重复框，或根据目标尺寸过滤误检的极小/极大框

内容的提问来源于stack exchange，提问作者Nikhil K Shaji