1. 环境准备与项目部署
第一次接触SAM自动标注工具时,我被它强大的零样本分割能力震撼到了。这个由Meta开源的Segment Anything Model(SAM)确实改变了传统标注工作的游戏规则。下面我就带大家从零开始搭建整套工具链,过程中会分享我踩过的坑和优化技巧。
先说说硬件需求。虽然官方推荐使用GPU加速,但实测在CPU环境下也能跑通整个流程,只是速度会慢3-5倍。我的测试环境是Ubuntu 20.04系统,Python 3.8环境,配RTX 3060显卡。如果你只有CPU,建议准备些咖啡耐心等待。
部署的第一步是克隆两个关键仓库:
git clone https://github.com/zhouayi/SAM-Tool.git git clone https://github.com/facebookresearch/segment-anything.git安装依赖时有个小技巧:先进入segment-anything目录执行可编辑安装,这样能避免后续路径问题:
cd segment-anything pip install -e .模型文件建议选择默认的sam_vit_h_4b8939.pth,这个2.4GB的模型在精度和速度上比较均衡。下载后记得校验MD5值,我就遇到过模型文件损坏导致后续报错的情况。
2. 数据准备与Embedding提取
数据组织是容易被忽视但极其重要的一环。建议采用这样的目录结构:
dataset/ ├── images/ │ ├── img1.jpg │ └── img2.png └── embeddings/ # 空目录图片格式支持JPG/PNG等常见格式,但要注意两点:一是文件名不要包含中文或特殊字符,二是建议统一尺寸。虽然SAM支持任意尺寸输入,但统一尺寸可以简化后续处理。
提取Embedding是核心步骤之一,这里有个性能优化点:
python helpers/extract_embeddings.py \ --checkpoint-path sam_vit_h_4b8939.pth \ --dataset-folder ./dataset \ --device cuda:0 # 使用指定GPU如果遇到CUDA内存不足的问题,可以添加--batch-size参数调小批次大小。提取完成后,每个图片会生成对应的.npy文件存放在embeddings目录,这些预处理结果能大幅提升后续标注效率。
3. 模型转换关键技巧
将PyTorch模型转为ONNX格式时,最容易踩的坑就是动态尺寸问题。原始代码生成的ONNX模型是固定尺寸的,这意味着:
- 如果训练图片尺寸不一,需要为每种尺寸单独导出模型
- 导出时orig-im-size参数必须与实际图片尺寸严格一致
转换命令示例:
python helpers/generate_onnx.py \ --checkpoint-path sam_vit_h_4b8939.pth \ --onnx-model-path ./sam_onnx_1080x1920.onnx \ --orig-im-size 1080 1920我遇到过opset版本不兼容的问题,解决方案是确保环境满足:
- PyTorch ≥ 1.13
- ONNX ≥ 1.13
- ONNX Runtime ≥ 1.14
如果导出失败,可以尝试指定opset版本:
python helpers/generate_onnx.py ... --opset-version 144. 交互标注实战技巧
终于到了最激动人心的标注环节!启动标注工具前,记得把ONNX模型复制到SAM-Tool目录:
cp sam_onnx.onnx ../SAM-Tool/启动标注界面时,类别参数很关键:
python segment_anything_annotator.py \ --onnx-model-path sam_onnx.onnx \ --dataset-path ./dataset \ --categories person,car,tree # 用英文逗号分隔实际操作时我发现这些技巧特别实用:
- 左键点击物体中心区域能得到最佳分割效果
- 按住Shift+左键拖动可以添加矩形提示框
- 右键点击误分割区域可以擦除错误部分
- 按K/L键实时调整掩码透明度方便观察
标注效率提升秘籍:
- 先粗略标注全部图片(按D键快速跳转)
- 第二遍集中优化有问题的标注
- 对相似物体使用复制标注功能(Ctrl+C/Ctrl+V)
5. 标注结果后处理
生成的COCO格式annotations.json包含完整标注信息。如果需要可视化检查:
python cocoviewer.py -i ./dataset -a ./dataset/annotations.json转换为VOC格式时,我优化过的脚本比原始版本更健壮:
def coco_to_voc(coco_json, output_dir): os.makedirs(output_dir, exist_ok=True) with open(coco_json) as f: data = json.load(f) for img in data['images']: xml_root = ET.Element('annotation') # 添加基础信息... for ann in filter(lambda x: x['image_id']==img['id'], data['annotations']): obj = ET.SubElement(xml_root, 'object') # 转换bbox坐标... ET.ElementTree(xml_root).write( os.path.join(output_dir, f"{os.path.splitext(img['file_name'])[0]}.xml") )对于大规模数据集,建议使用多进程加速转换:
from multiprocessing import Pool with Pool(8) as p: # 8进程并行 p.map(coco_to_voc, image_batches)6. 性能优化与疑难解答
经过多个项目实践,我总结出这些优化方案:
CPU模式加速技巧:
- 设置OMP_NUM_THREADS环境变量
- 在extract_embeddings.py中启用--use-flash-attention
- 使用ONNX Runtime的性能模式
常见错误解决方案:
"ValueError: Unsupported ONNX opset version"
- 升级onnx和onnxruntime包
- 或降低opset版本至14
"CUDA out of memory"
- 减小--batch-size参数值
- 尝试--precision fp16混合精度
标注结果不准确
- 检查图片是否过度压缩
- 尝试不同的提示点位置
- 调整--pred-iou-thresh参数
对于专业标注团队,我建议开发这些扩展功能:
- 自定义快捷键配置
- 批量修正工具
- 质量检查自动化脚本
- 与Labelme等工具的互转换
这套工具链在实际项目中表现惊人,原本需要2周的人工标注工作,现在2-3天就能完成,且质量更稳定。特别是在医疗影像和遥感图像这些专业领域,SAM展现出了超越传统方法的泛化能力。