一、项目背景

在目标检测任务中，数据标注是一项耗时且繁琐的工作。当我在有一个训练好的YOLO模型后，想要更快地迭代后续的模型版本，于是我在网上找自动化标注的资料和教程，有很多资料都是把已有模型接入到数据标注平台，需要各种配置安装，还有很多限制和要求，于是我就想用python脚本来实现自动标注的过程，简单又高效，大大提高标注效率。这篇文章介绍如何利用YOLO模型自动标注images文件夹中的图片，并输出多种格式的标注结果。

 二、文件夹结构

project/

├── model/

│   └── best.pt           # 训练好的YOLO模型

├── images/               # 待标注的原始图片,jpg或png格式

│   └── .......jpg        

│   └── .......jpg        
......

文件夹结构仅供参考，如果修改文件夹结构，运行的代码也要做相应修改

 三、核心代码

from ultralytics import YOLO
from pathlib import Path

model_path = Path(__file__).parent / "model" / "best.pt"
images_dir = Path(__file__).parent / "images"
output_dir = Path(__file__).parent / "output"

output_dir.mkdir(exist_ok=True)

model = YOLO(str(model_path))

image_files = list(images_dir.glob("*.jpg")) + list(images_dir.glob("*.png"))

print(f"找到 {len(image_files)} 张图片")

for img_path in image_files:
    results = model(str(img_path))
    result = results[0]

    detections = []
    if result.boxes is not None:
        for box in result.boxes:
            cls_id = int(box.cls[0])
            conf = float(box.conf[0])
            bbox = box.xyxy[0].tolist()
            detections.append({
                'class_id': cls_id,
                'confidence': conf,
                'bbox': bbox
            })

    # ========== 扩展功能1: 添加置信度过滤 ==========
    # conf_threshold = 0.5
    # detections = [d for d in detections if d['confidence'] >= conf_threshold]

    txt_path = output_dir / f"{img_path.stem}.txt"
    with open(txt_path, 'w') as f:
        for det in detections:
            f.write(f"{det['class_id']} {det['confidence']:.4f} {det['bbox'][0]:.2f} {det['bbox'][1]:.2f} {det['bbox'][2]:.2f} {det['bbox'][3]:.2f}\n")

    # ========== 扩展功能2: 生成可视化标注图 ==========
    # annotated_img = result.plot()
    # result.save(f"annotated_{img_path.name}")

    # ========== 扩展功能3: 输出VOC XML格式 ==========
    # from PIL import Image
    # img = Image.open(str(img_path))
    # img_width, img_height = img.size
    # generate_voc_xml(img_path.name, img_width, img_height, detections, model.names)

    print(f"处理完成: {img_path.name}")

print(f"标注完成！结果保存在: {output_dir}")

 

四、使用方法

 4.1 环境准备

pip install ultralytics

 4.2 操作步骤

1.按照文件夹结构放好对应的文件。

2.按照自己的电脑情况微调代码，比如模型及文件夹的名称和路径。

3.执行核心代码进行自动标注，并等待标注结果。

4.标注后的数据集可以做相应的格式调整后导入到数据标注平台进行人工调整与验证，也可以用于后续新模型的训练。

 五、输出格式说明

该脚本默认输出为YOLO TXT格式，运行结束后会在output文件夹中生成标签文件

每张图片对应一个 txt标签文件，每行表示一个检测框：

```

class_id confidence x1 y1 x2 y2

```


**参数说明：**

- `class_id`: 类别索引（从0开始）

- `confidence`: 检测置信度（0-1）

- `x1, y1`: 检测框左上角坐标

- `x2, y2`: 检测框右下角坐标


**示例：**

```

0 0.9523 120.00 350.00 450.00 720.00

1 0.8912 200.00 100.00 350.00 280.00

```

六、注意事项

1. 模型路径：确保模型文件存在于 `model/best.pt`，注意名称不要错了

2. 图片格式：支持 jpg 和 png 格式

3. 批量处理：脚本会自动处理文件夹中的所有图片

4. 空标注：如果图片中没有检测到目标，仍会生成空的 txt 文件

七、扩展功能建议

扩展功能代码已在上面的核心代码中用 `========== 扩展功能N ==========` 标注了具体位置，可以直接取消注释使用。

如果对你有帮助，欢迎在点赞收藏在评论区留言交流！