使用Segment Anything(SAM)模型进行自动标注
详细教程,使用Segment Anything(SAM)模型当工具进行自己数据的自动标注。
1.下载项目
项目1:https://github.com/zhouayi/SAM-Tool
项目2:https://github.com/facebookresearch/segment-anything
git clone https://github.com/zhouayi/SAM-Tool.git
git clone https://github.com/facebookresearch/segment-anything.git
cd segment-anything
pip install -e .
下载SAM
模型:https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth
2.把数据放置在<dataset_path>/images/*
这样的路径中,并创建空文件夹<dataset_path>/embeddings
3.将项目1中的helpers
文件夹复制到项目2的主目录下
3.1 运行extract_embeddings.py
文件来提取图片的embedding
# cd到项目2的主目录下
python helpers\extract_embeddings.py --checkpoint-path sam_vit_h_4b8939.pth --dataset-folder <dataset_path> --device cpu
checkpoint-path
:上面下载好的SAM
模型路径dataset-folder
:数据路径device
:默认cuda
,没有GPU
用cpu
也行的,就是速度挺慢的
运行完毕后,<dataset_path>/embeddings
下会生成相应的npy文件
3.2 运行generate_onnx.py
将pth
文件转换为onnx
模型文件
# cd到项目2的主目录下
python helpers\generate_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 <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
格式,保存在<dataset_path>/annotations.json
。
5. 查看全部标注结果
运行cocoviewer.py
查看全部的标注结果
# cd到项目1的主目录下
python cocoviewer.py -i <dataset_path> -a <dataset_path>\annotations.json
6.将保存的json格式转换为voc格式
该工具保存的标注文件格式为COCO标准格式,如有需求可自行编写标注文件格式转换脚本,下面提供转换为VOC格式的脚本例子
import json
import xml.etree.ElementTree as ET
import os
jsonPath = "F:/vsCode/segment-anything/CatDog/annotations.json"
vocPath = "F:/vsCode/segment-anything/CatDog/VOC"
with open(jsonPath, 'r') as f:
data = json.load(f)
info = data["info"]
images = data["images"]
annotations = data["annotations"]
categories = data["categories"]
# 对每个图像处理
for img_data in images:
# 创建 VOC XML 文件
xml_file = ET.Element('annotation')
ET.SubElement(xml_file, 'folder').text = 'VOC'
ET.SubElement(xml_file, 'filename').text = os.path.basename(img_data["file_name"])
source = ET.SubElement(xml_file, 'source')
ET.SubElement(source, 'database').text = 'My Database'
ET.SubElement(source, 'annotation').text = 'COCO'
ET.SubElement(source, 'image').text = 'flickr'
size = ET.SubElement(xml_file, 'size')
ET.SubElement(size, 'width').text = str(img_data['width'])
ET.SubElement(size, 'height').text = str(img_data['height'])
ET.SubElement(size, 'depth').text = '3'
ET.SubElement(xml_file, 'segmented').text = '0'
# 查找该图像的所有标注框
bbox_list = []
category_ids = []
for ann_data in annotations:
if ann_data['image_id'] == img_data['id']:
bbox = ann_data['bbox']
bbox_list.append(bbox)
category_ids.append(ann_data['category_id'])
# 对每个标注框处理
for i in range(len(bbox_list)):
bbox = bbox_list[i]
category_id = category_ids[i]
# 转换 COCO 格式到 VOC 格式
x_min = bbox[0]
y_min = bbox[1]
x_max = bbox[0] + bbox[2]
y_max = bbox[1] + bbox[3]
class_name = categories[category_id]['name']
# 创建 VOC XML 标注
obj = ET.SubElement(xml_file, 'object')
ET.SubElement(obj, 'name').text = class_name
ET.SubElement(obj, 'pose').text = 'Unspecified'
ET.SubElement(obj, 'truncated').text = '0'
ET.SubElement(obj, 'difficult').text = '0'
bndbox = ET.SubElement(obj, 'bndbox')
ET.SubElement(bndbox, 'xmin').text = str(int(x_min))
ET.SubElement(bndbox, 'ymin').text = str(int(y_min))
ET.SubElement(bndbox, 'xmax').text = str(int(x_max))
ET.SubElement(bndbox, 'ymax').text = str(int(y_max))
# 将 XML 文件保存到 VOC 目标文件夹中
xml_str = ET.tostring(xml_file)
with open(os.path.join(vocPath, os.path.basename(img_data["file_name"]).replace('.jpg', '.xml')), 'wb') as f:
f.write(xml_str)
VOC格式:
7.其他
- 修改标注框线条的宽度的代码位置
# salt/displat_utils.py
class DisplayUtils:
def __init__(self):
self.transparency = 0.65 # 默认的掩码透明度
self.box_width = 2 # 默认的边界框线条宽度
- 修改标注文本的格式的代码位置
# salt/displat_utils.py
def 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
为武汉地区的开发者提供学习、交流和合作的平台。社区聚集了众多技术爱好者和专业人士,涵盖了多个领域,包括人工智能、大数据、云计算、区块链等。社区定期举办技术分享、培训和活动,为开发者提供更多的学习和交流机会。
更多推荐
所有评论(0)