1.下载项目
项目1:https://github.com/zhouayi/SAM-Tool
项目2:https://github.com/facebookresearch/segment-anything
git clone https://github.com/zhouayi/SAM-Tool.gitgit clone https://github.com/facebookresearch/segment-anything.gitcd segment-anythingpip install -e .
下载SAM模型:https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth
2.把数据放置在/images/*这样的路径中,并创建空文件夹/embeddings
3.将项目1中的helpers文件夹复制到项目2的主目录下
3.1 运行extract_embeddings.py文件来提取图片的embedding
# cd到项目2的主目录下python helpersextract_embeddings.py --checkpoint-path sam_vit_h_4b8939.pth --dataset-folder--device cpu
- checkpoint-path:上面下载好的SAM模型路径
- dataset-folder:数据路径
- device:默认cuda,没有GPU用cpu也行的,就是速度挺慢的
运行完毕后,
3.2 运行generate_onnx.py将pth文件转换为onnx模型文件
# cd到项目2的主目录下python helpersgenerate_onnx.py --checkpoint-path sam_vit_h_4b8939.pth --onnx-model-path ./sam_onnx.onnx --orig-im-size 1080 1920
- checkpoint-path:同样的SAM模型路径
- onnx-model-path:得到的onnx模型保存路径
- orig-im-size:数据中图片的尺寸大小(height, width)
【注意:提供给的代码转换得到的onnx模型并不支持动态输入大小,所以如果你的数据集中图片尺寸不一,那么可选方案是以不同的orig-im-size参数导出不同的onnx模型供后续使用】
这里可能会报错ValueError: Unsupported ONNX opset version: 15,原因是因为pytorch或者onnx版本太低了,解决方法是安装更高版本的pytorch和onnx,可供参考,我的环境是
onnx == 1.13.1
onnxruntime == 1.14.1
torch == 1.13.1
torchaudio== 0.13.1
torchvision==0.14.1
我同样也尝试了另一种方法:将opset-version数值设低一点,但是会报其他的错误
4.将生成的sam_onnx.onnx模型复制到项目1的主目录下,运行segment_anything_annotator.py进行标注
# cd到项目1的主目录下python segment_anything_annotator.py --onnx-model-path sam_onnx.onnx --dataset-path--categories cat,dog
- onnx-model-path:导出的onnx模型路径
- dataset-path:数据路径
- categories:数据集的类别(每个类别以,分割,不要有空格)
在对象位置处点击鼠标左键为增加掩码,点击右键为去掉该位置掩码。
其他使用快捷键有:
Esc:退出app | a:前一张图片 | d:下一张图片 |
---|---|---|
k:调低透明度 | l:调高透明度 | n:添加对象 |
r:重置 | Ctrl+s:保存 | Ctrl+z:撤销上一个对象 |
最后生成的标注文件为coco格式,保存在
5. 查看全部标注结果
运行cocoviewer.py查看全部的标注结果
# cd到项目1的主目录下python cocoviewer.py -i-a annotations.json
6.其他
- 修改标注框线条的宽度的代码位置
# salt/displat_utils.pyclass DisplayUtils: def __init__(self): self.transparency = 0.65 # 默认的掩码透明度 self.box_width = 2 # 默认的边界框线条宽度
- 修改标注文本的格式的代码位置
# salt/displat_utils.pydef draw_box_on_image(self, image, categories, ann, color): x, y, w, h = ann["bbox"] x, y, w, h = int(x), int(y), int(w), int(h) image = cv2.rectangle(image, (x, y), (x + w, y + h), color, self.box_width) text = '{} {}'.format(ann["id"],categories[ann["category_id"]]) txt_color = (0, 0, 0) if np.mean(color) > 127 else (255, 255, 255) font = cv2.FONT_HERSHEY_SIMPLEX txt_size = cv2.getTextSize(text, font, 1.5, 1)[0] cv2.rectangle(image, (x, y + 1), (x + txt_size[0] + 1, y + int(1.5*txt_size[1])), color, -1) cv2.putText(image, text, (x, y + txt_size[1]), font, 1.5, txt_color, thickness=5) return image
- 2023.04.14新增加撤销上一个标注对象功能,快捷键Ctrl+z