基于 YOLO + SAM 的自动图像标注实战

用检测模型做“找对象”，用分割模型做“画轮廓”，一套脚本把数据标注从体力活升级成自动化流水线。

一、背景与动机

在计算机视觉项目中，高质量标注数据往往比模型结构本身更昂贵。尤其是：

• 目标类别多
• 图像数量大
• 需要实例级 / 分割级标注

手工标注不仅慢，而且极易成为项目瓶颈。

一个非常自然的想法是：

能不能让模型帮我们先标一版？

本文介绍一种在工程中非常实用的方案：

• YOLO：负责目标检测（给出 bbox + 类别）
• SAM（Segment Anything Model）：基于 bbox 进一步做精细分割

最终目标：

输入一堆图片 → 自动生成可用于训练的标注文件（YOLO 格式）

二、整体流程设计

整体流程非常清晰：

Images
  ↓
YOLO 检测（bbox + class）
  ↓
SAM 分割（基于 bbox）
  ↓
生成 txt 标注文件

核心思想

• YOLO 擅长 “找在哪”
• SAM 擅长 “抠得准”
• 用 YOLO 的 bbox 作为 prompt，可以显著提升 SAM 的稳定性和速度

三、代码结构总览

脚本主要由 4 个部分组成：

1. 模型加载（YOLO + SAM）
2. 单张图片处理逻辑
3. 批量数据处理
4. 命令行接口设计

下面逐一拆解。

四、模型加载：健壮第一

def load_models(yolo_model_path: str, sam_model_path: str):
    yolo_model = YOLO(yolo_model_path)
    sam_model = SAM(sam_model_path)
    return yolo_model, sam_model, yolo_model.names

设计要点

• 显式异常捕获：模型路径错误时能快速定位问题
• 同时返回 obj_names，用于生成 classes.txt

这是一个小细节，但在数据流水线中非常重要。

五、单张图片处理逻辑（核心）

1️⃣ 读取图像

img = cv2.imread(str(img_path))

• 无法读取直接跳过
• 防止坏图导致整个流程中断

2️⃣ YOLO 目标检测

results = yolo_model.predict(img, verbose=