ARTICLE DETAIL

资讯详情

深耕郑州网站建设与运营推广的一线实战洞察。

YOLO+SAM自动图像标注:从检测框到像素级掩码的工程实践

YOLO+SAM自动图像标注:从检测框到像素级掩码的工程实践 1. 自动图像标注的底层逻辑与方案选型做过目标检测项目的人都有一个共识模型训练本身花的时间往往远不如标注数据花的时间多。一个中等规模的数据集几千张图纯手工拉框一个人干一周是常态标注质量还参差不齐。我最早做工业质检项目的时候三千张缺陷图三个人标了整整五天最后一致性检查还发现百分之十几的框有偏差。所以当 YOLO 和 SAM 这两个东西凑到一起的时候我第一反应就是——这事有搞头。这个项目的核心思路其实不复杂用 YOLO 做粗定位用 SAM 做精分割两者串起来把“检测框”自动升级成“像素级掩码”再反向生成标注文件。你只需要标一小部分数据训练一个 YOLO 检测器剩下的图让模型自己跑跑完再用 SAM 把框内的目标抠出来最后导出成 LabelImg 或者 CVAT 能直接读的格式。整个流程下来标注效率提升五到十倍是保守估计。为什么选 YOLO 而不是别的检测器原因很直接部署简单、推理快、小样本微调效果好。你拿一个预训练的 YOLOv8n 或者 YOLOv11n冻结主干只训检测头几百张图就能出一个可用的基线。换成 DETR 系列收敛慢小数据集上容易过拟合调参成本高。而 SAM 的价值在于它的零样本分割能力——你给它一个框或者一个点它就能把目标轮廓抠出来不需要针对每个类别重新训练分割头。这两者结合等于用最低的训练成本换最高的标注自动化程度。这里要澄清一个常见误解SAM 不是用来做检测的。它的输出是掩码不是类别。所以你不能指望 SAM 告诉你“这是猫还是狗”它只负责“把这一坨像素圈出来”。类别信息必须由 YOLO 提供。这也是为什么这个方案是“YOLO SAM”而不是“SAM 单独搞定一切”。分工明确各干各擅长的事。适合谁来参考这套方案如果你正在做目标检测项目手头有几百到几千张图需要标注预算有限不想买商业标注服务同时又有一定的 Python 基础和显卡资源那这套流程就是为你准备的。哪怕你之前没接触过 SAM只要会跑 YOLO 推理跟着走一遍就能上手。2. 环境搭建与核心依赖的版本坑2.1 硬件与基础环境的最低要求先说硬门槛。YOLO 推理本身很轻GTX 1060 6G 都能跑。但 SAM 不一样ViT-H 版本的模型文件 2.4G推理时显存占用在 4G 到 6G 之间浮动取决于输入分辨率。如果你用的是 SAM 的 ViT-B 版本模型小很多显存 2G 左右就够但分割边缘的精细度会下降。我的建议是显卡至少 8G 显存这样 YOLO 和 SAM 可以同时加载不用反复切换模型批量处理时效率高很多。操作系统层面Windows 和 Linux 都行。Windows 上用 Anaconda 建虚拟环境最省事Linux 上直接用 venv 或者 conda 都可以。Python 版本锁定在 3.9 到 3.10 之间太新的版本有些依赖包还没跟上太老的版本 PyTorch 又不支持。CUDA 版本跟着 PyTorch 走目前 PyTorch 2.1 以上配 CUDA 11.8 或 12.1 都比较稳。2.2 依赖安装的先后顺序与常见报错安装顺序很重要搞错了就是无尽的版本冲突。我的习惯是先建虚拟环境conda create -n auto_label python3.10装 PyTorch去官网查对应 CUDA 版本的安装命令别直接pip install torch那样装的是 CPU 版装 ultralyticspip install ultralytics这是 YOLO 的官方库更新很勤装 segment-anythingpip install githttps://github.com/facebookresearch/segment-anything.git装辅助库opencv-python、pillow、numpy、tqdm、pyyaml这里有个坑我踩过ultralytics 和 segment-anything 对 numpy 版本的要求可能冲突。ultralytics 新版要求 numpy1.23而某些旧版的 segment-anything 依赖 numpy1.24。解决办法是先装 ultralytics再装 segment-anything如果报错就手动指定 numpy 版本pip install numpy1.23.5。实测这个版本两边都能兼容。另一个常见问题是 SAM 的模型权重下载。官方提供了三个版本ViT-B375M、ViT-L1.2G、ViT-H2.4G。国内下载速度可能很慢建议提前用下载工具拉下来放到指定目录代码里直接加载本地路径。模型文件放哪无所谓只要路径写对就行。注意不要用 pip 直接装 opencv-python-headless 和 opencv-python 两个包会冲突。只装 opencv-python 就够了headless 版本没有 GUI 功能调试时不方便。2.3 验证环境是否可用的最小测试装完之后别急着跑全流程先做个最小验证。用 YOLO 跑一张测试图确认能出框再用 SAM 加载模型给一个手动框确认能出掩码。这两步都通了再往下走。我见过太多人环境没验就直接上批量脚本结果跑到一半报错浪费一晚上。from ultralytics import YOLO from segment_anything import sam_model_registry, SamPredictor import cv2 # YOLO 验证 model YOLO(yolov8n.pt) results model(test.jpg) print(检测到, len(results[0].boxes), 个目标) # SAM 验证 sam sam_model_registry[vit_b](checkpointsam_vit_b_01ec64.pth) sam.to(cuda) predictor SamPredictor(sam) img cv2.imread(test.jpg) predictor.set_image(img) masks, scores, _ predictor.predict(boxresults[0].boxes[0].xyxy.cpu().numpy()[0]) print(掩码形状, masks.shape)这段代码跑通说明基础环境没问题。如果 SAM 那步报显存不足换成 ViT-B 版本再试。3. YOLO 检测器的训练与调优细节3.1 数据集准备与标注策略自动标注的前提是你得先有一个能用的检测器。这个检测器不需要很完美但必须能稳定召回目标。我的做法是从全部数据里随机抽 10% 到 20% 做人工标注训练一个基线 YOLO。剩下的 80% 到 90% 交给自动流程。标注工具用 LabelImg 就行导出 YOLO 格式的 txt 文件。每张图对应一个 txt每行是类别id 中心x 中心y 宽 高坐标都归一化到 0 到 1 之间。这里有个细节标注框要尽量贴紧目标边缘但不要切掉目标。因为后面 SAM 会以这个框为提示做分割框太松会导致 SAM 把背景也抠进去框太紧又可能让 SAM 丢失边缘细节。类别数量少的时候每个类至少标 100 到 150 个实例。类别多的话优先保证长尾类别的样本量因为 YOLO 对长尾类别天然不友好。如果某个类实在样本太少可以考虑用数据增强先扩充一波或者用其他类的预训练权重做迁移。3.2 YOLO 训练参数的选择逻辑训练 YOLOv8 或者 YOLOv11 的时候我一般用这样的配置参数推荐值说明epochs100-150小数据集容易过拟合早停 patience20imgsz640和推理尺寸保持一致batch8-16根据显存调整显存不够就降lr00.001默认值小数据集可以降到 0.0005freeze10冻结主干前10层加速收敛augmentTrue开启 mosaic 和 mixup为什么要冻结主干因为你的数据集小全量微调容易把预训练学到的通用特征破坏掉。冻结前 10 层只训检测头和后几层收敛更快泛化也更好。实测下来冻结训练比不冻结在 500 张图的数据集上 mAP 高 3 到 5 个点。训练完成后看验证集上的 mAP50 和 mAP50-95。mAP50 到 0.85 以上基本就够用了。如果低于 0.7要么加数据要么检查标注质量。别在检测器上死磕太久因为后面 SAM 会弥补一部分定位精度检测器只要能把目标框住就行。3.3 推理阶段的置信度与 NMS 设置自动标注时YOLO 的推理参数和正常部署不太一样。正常部署追求高精度置信度阈值可以设 0.5。但自动标注追求高召回宁可多框几个也不能漏。所以我把conf设到 0.25 到 0.3iou设到 0.5 到 0.6。results model(img, conf0.25, iou0.5, verboseFalse)这样设置的好处是一些遮挡严重或者模糊的目标也能被框出来。代价是可能有一些误检但误检的框后面可以用 SAM 的分割质量分数过滤掉——如果 SAM 对某个框给出的掩码置信度很低说明这个框可能本身就不靠谱直接丢弃。实操心得批量推理时把verboseFalse加上不然控制台会刷屏影响你看进度。另外streamTrue可以节省内存适合处理大量图片。4. SAM 分割的提示工程与掩码后处理4.1 用检测框作为提示的正确姿势SAM 支持三种提示方式点、框、掩码。自动标注场景下YOLO 给的是框所以直接用框提示最自然。但这里有个细节SAM 的框提示格式是 xyxy不是 YOLO 的 xywh。YOLO 输出的是归一化的中心点加宽高需要先转换成绝对坐标的左上右下。boxes results[0].boxes.xyxy.cpu().numpy() # 已经是 xyxy 绝对坐标 for box in boxes: masks, scores, _ predictor.predict(boxbox, multimask_outputTrue) best_mask masks[np.argmax(scores)]multimask_outputTrue会让 SAM 输出三个候选掩码分别对应不同粒度。选分数最高的那个通常是最贴合目标轮廓的。如果目标特别小或者边缘特别复杂可以试试multimask_outputFalse只出一个掩码有时候反而更稳。4.2 掩码的过滤与形态学处理SAM 出来的掩码不是每个都能用。有些掩码会溢出到背景有些会漏掉目标的一部分。我一般做三层过滤第一层面积过滤。掩码面积和检测框面积的比值如果小于 0.3说明 SAM 只抠出了一小部分大概率是失败的。如果大于 1.2说明掩码溢出了框也不正常。正常范围在 0.5 到 1.0 之间。第二层置信度过滤。SAM 返回的 score 低于 0.7 的直接丢弃。这个阈值可以根据你的数据调数据干净就设高一点数据杂乱就设低一点。第三层形态学处理。用开运算去掉小噪点用闭运算填补内部空洞。OpenCV 的morphologyEx一行搞定。kernel np.ones((3,3), np.uint8) mask cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel) mask cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel)经过这三层处理剩下的掩码质量就相当可观了。实测在工业缺陷数据集上自动标注的掩码和人工标注的 IoU 能达到 0.85 以上。4.3 从掩码到标注文件的转换得到掩码之后要转成标注工具能读的格式。如果你只需要检测框那直接把 YOLO 的框导出来就行。但如果你要做实例分割就需要把掩码转成多边形或者 RLE 编码。转多边形用 OpenCV 的findContourscontours, _ cv2.findContours(mask.astype(np.uint8), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) for contour in contours: epsilon 0.001 * cv2.arcLength(contour, True) approx cv2.approxPolyDP(contour, epsilon, True) polygon approx.reshape(-1, 2)approxPolyDP做多边形逼近epsilon 控制精度。太小了点数多标注文件大太大了轮廓失真。0.001 到 0.002 倍周长是比较平衡的值。转出来的多边形可以存成 COCO 格式的 json也可以存成 YOLO 分割格式的 txt。YOLO 分割格式每行是类别id x1 y1 x2 y2 ...坐标同样归一化。LabelImg 不直接支持分割格式但 CVAT 和 LabelMe 可以导入。5. 批量处理流程与性能优化5.1 批量推理的脚本结构单张图跑通了接下来就是批量。批量脚本的核心逻辑是遍历图片目录YOLO 推理SAM 分割过滤导出。但直接串行跑会很慢因为 SAM 的set_image这一步很耗时每张图都要重新编码。优化思路是YOLO 批量推理SAM 逐张分割。YOLO 可以用streamTrue一次处理多张SAM 因为要维护图像编码只能一张一张来。但可以把 YOLO 的结果先存下来再统一跑 SAM这样两个阶段解耦方便中断续跑。# 阶段一YOLO 批量推理 all_boxes {} for result in model(source_dir, streamTrue, conf0.25): img_name Path(result.path).name all_boxes[img_name] result.boxes.xyxy.cpu().numpy() # 阶段二SAM 逐张分割 for img_name, boxes in tqdm(all_boxes.items()): img cv2.imread(str(source_dir / img_name)) predictor.set_image(img) for box in boxes: masks, scores, _ predictor.predict(boxbox, multimask_outputTrue) # 过滤和保存逻辑这样写的好处是如果 SAM 跑到一半崩了YOLO 的结果还在不用重跑。5.2 显存管理与批大小调优SAM 的显存占用主要来自图像编码器。输入分辨率越高显存越大。默认 SAM 会把图像缩放到 1024 的长边如果你的原图是 4000 像素宽缩放后信息损失很多小目标的分割质量会下降。解决办法是分块推理把大图切成 1024 乘 1024 的小块每块单独跑 SAM最后把掩码拼回去。这样显存占用可控小目标也能保住细节。代价是边缘处可能有拼接痕迹需要用重叠区域做融合。如果显存实在紧张可以降低 SAM 的输入分辨率。SamPredictor有个set_image方法内部会做缩放。你也可以手动把图缩小再传进去但要注意把框的坐标同步缩放。5.3 处理速度的实测数据与瓶颈分析我在一台 RTX 3060 12G 的机器上做过测试一千张 1920 乘 1080 的图片阶段耗时说明YOLO 推理约 3 分钟batch16FP16SAM 编码约 25 分钟ViT-B逐张SAM 解码约 8 分钟每张图平均 5 个目标后处理与导出约 2 分钟形态学加轮廓提取合计约 38 分钟平均每张 2.3 秒瓶颈很明显在 SAM 的图像编码。ViT-H 会更慢大概慢 2 到 3 倍。如果对速度要求高可以用 ViT-B或者用 ONNX Runtime 加速 SAM 的编码器。我试过用 TensorRT 部署 SAM编码速度能提升 40% 左右但转换过程比较折腾适合有部署经验的人。6. 常见问题排查与避坑指南6.1 检测框漏检导致的标注缺失自动标注最怕的就是漏检。YOLO 没框到的地方SAM 根本不会去分割最后标注文件里就少了这些目标。漏检的原因通常有三个置信度阈值设太高、目标太小、类别样本不均衡。排查方法把conf降到 0.1重新跑一遍看漏检的目标能不能被框出来。如果能说明是阈值问题适当降低阈值。如果还是不行说明检测器本身没学好需要补标注数据重新训练。对于小目标可以把imgsz从 640 提到 1280推理时放大输入小目标的召回会明显改善。避坑技巧自动标注完成后随机抽 50 张图做人工复核统计漏检率和误检率。如果漏检超过 5%这批自动标注就不能直接用得回去补数据。6.2 SAM 掩码溢出与欠分割的处理SAM 掩码溢出通常发生在目标和背景颜色接近的时候。比如白墙上的白色物体SAM 很容易把整面墙都抠进去。解决办法是收紧提示框把 YOLO 的框往里缩 5% 到 10%给 SAM 更明确的边界信号。欠分割则是掩码只覆盖了目标的一部分。这通常是因为目标内部有纹理断裂或者目标被遮挡。可以试试用多个点提示代替框提示在目标内部均匀撒几个正点在背景撒几个负点引导 SAM 更准确地分割。SamPredictor的predict方法支持同时传point_coords和point_labels。point_coords np.array([[x1,y1],[x2,y2],[x3,y3]]) point_labels np.array([1,1,0]) # 1是前景0是背景 masks, scores, _ predictor.predict(point_coordspoint_coords, point_labelspoint_labels, boxbox)6.3 标注格式转换中的坐标对齐问题从掩码转多边形再从多边形转 YOLO 格式中间涉及多次坐标变换很容易出错。最常见的错误是归一化时用错了图像尺寸。YOLO 格式要求坐标除以原图的宽和高但如果你在缩放后的图上做的分割就得先还原到原图尺寸再归一化。我的做法是所有坐标变换都在原图坐标系下完成只在最后导出时做一次归一化。中间过程用绝对坐标避免反复缩放带来的精度损失。另外多边形的点顺序要统一顺时针或者逆时针都行但不能混。混了会导致某些标注工具解析出错。6.4 常见问题速查表问题现象可能原因解决方法SAM 报显存不足模型太大或图太大换 ViT-B或分块推理掩码全是背景框提示不准确收紧框或加点提示标注文件打不开格式或坐标错误检查归一化和点顺序YOLO 漏检严重阈值高或数据少降阈值补标注数据处理速度太慢SAM 编码瓶颈换小模型或 TensorRT 加速掩码边缘毛刺分辨率不够提高输入分辨率或后处理平滑7. 自动标注结果的质量评估与迭代自动标注不是一锤子买卖跑完一轮之后必须评估质量然后针对性迭代。我的评估流程分三步第一步抽样人工复核。从自动标注结果里随机抽 10% 的图人工检查框和掩码的准确性。记录漏检数、误检数、掩码 IoU 低于 0.7 的数量。第二步计算自动指标。如果有部分人工标注的 ground truth可以直接算 mAP 和 mask IoU。没有的话就用 SAM 的置信度分数做代理指标分数分布偏低说明整体质量有问题。第三步迭代策略。如果漏检多就补标注漏检的类别重新训 YOLO。如果掩码质量差就调 SAM 的提示策略或者换更大的 SAM 模型。如果误检多就提高 YOLO 的置信度阈值或者加一个分类器做二次过滤。我自己的项目里第一轮自动标注通常能达到 80% 左右的可用率剩下 20% 需要人工修正。第二轮把修正后的数据加入训练集重新训 YOLO可用率能提到 90% 以上。第三轮基本就稳定了人工只需要做最终抽检。这套流程最大的价值不是完全替代人工而是把人工从重复劳动里解放出来只做最关键的审核和修正。对于标注预算有限、时间紧张的团队来说这是目前性价比最高的方案之一。
返回列表