Uso de SAM para realizar el etiquetado automático
Tabla de contenido
1. Descarga e instala
1.1 Descargar SAM-Tool y SAM
clon de git https://github.com/zhouayi/SAM-Tool.git
clon de git https://github.com/facebookresearch/segment-anything.git
1.2 Descargar archivo modelo SAM
wget https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth
1.3 instalar SAM
Ingrese a la carpeta SAM: cd segmento-cualquier cosa
Instale SAM: pip install -e .
2. Elementos de configuración
(1) Mueva la carpeta de ayuda en la herramienta SAM a la carpeta principal de SAM: cp -r helpers/ …/segment-anything/ (
2) Cree el formato de la carpeta de datos como dataset_path/images y dataset_path/embeddings (usado en imágenes para poner fotos)
3. Extraer información
(1) Extraiga la incrustación de imágenes: python helpers/extract_embeddings.py --checkpoint-path sam_vit_h_4b8939.pth --dataset-folder Su dirección de conjunto de datos --device cuda
(2) Después de ejecutar, habrá un archivo npy correspondiente debajo de incrustaciones
Cuarto, obtenga el archivo SAM onnx
(1) Primero, debe haber onnx y onnxruntime(-gpu) en el entorno
(2) El archivo para detectar imágenes (se puede ajustar el alto/ancho de las siguientes imágenes): python helpers/generate_onnx.py --checkpoint- ruta sam_vit_h_4b8939.pth --onnx -model-path ./sam_onnx.onnx --orig-im-size 720 1280
(3) Habrá un archivo sam_onnx.onnx correspondiente después de ejecutarlo, muévalo a la carpeta principal de la herramienta SAM : cp sam_onnx.onnx …/SAM-Herramienta/
5. Etiquetado
(1) La biblioteca opencv necesita eliminar opencv-python y usar la versión sin cabeza en su lugar: pip install opencv-python-headless
(2) El comando para llamar al programa de anotación es (puede cambiar la categoría de anotación más adelante, separada por comas) : python segment_anything_annotator.py --onnx-model-path sam_onnx.onnx --dataset-path /home/dxfcv/workspace/sunsirui/label/segment-anything/dataset --categories cement,metal,plastics (3) Empezar a etiquetar
después aparece el cuadro de etiqueta, primero haga clic en el objetivo, luego agregue el objeto, finalmente seleccione la categoría y luego la siguiente, guarde todo después de marcar, y finalmente el archivo de formato coco de anotaciones.json aparecerá en la carpeta de datos (4
) Vea el archivo de anotaciones: python cocoviewer.py -i .../segment-anything/dataset/ -a .../segment-anything/dataset/annotations.json
6. Conversión de formato
6.1 convertir json a voc
import json
import xml.etree.ElementTree as ET
import os
jsonPath = "/home/dxfcv/workspace/sunsirui/label/dataset/annotations.json"#改这个输入json
vocPath = "/home/dxfcv/workspace/sunsirui/label/dataset/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)
6.2 Convertir json a yolo
import os
import json
from tqdm import tqdm
import argparse
def convert(size, box):
'''
size: 图片的宽和高(w,h)
box格式: x,y,w,h
返回值:x_center/image_width y_center/image_height width/image_width height/image_height
'''
dw = 1. / (size[0])
dh = 1. / (size[1])
x = box[0] + box[2] / 2.0
y = box[1] + box[3] / 2.0
w = box[2]
h = box[3]
x = x * dw
w = w * dw
y = y * dh
h = h * dh
return (x, y, w, h)
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument('--json_file', default='/home/dxfcv/workspace/sunsirui/label/dataset/annotations.json',
type=str, help="coco file path")
parser.add_argument('--save_dir', default='/home/dxfcv/workspace/sunsirui/label/dataset/labels', type=str,
help="where to save .txt labels")
arg = parser.parse_args()
data = json.load(open(arg.json_file, 'r'))
# 如果存放txt文件夹不存在,则创建
if not os.path.exists(arg.save_dir):
os.makedirs(arg.save_dir)
id_map = {
}
# 解析目标类别,也就是 categories 字段,并将类别写入文件 classes.txt 中
with open(os.path.join(arg.save_dir, 'classes.txt'), 'w') as f:
for i, category in enumerate(data['categories']):
f.write(f"{
category['name']}\n")
id_map[category['id']] = i
for img in tqdm(data['images']):
# 解析 images 字段,分别取出图片文件名、图片的宽和高、图片id
# filename = os.path.join(arg.image_file , img["file_name"])
filename = img["file_name"].split('/')[1]
img_width = img["width"]
img_height = img["height"]
img_id = img["id"]
head, tail = os.path.splitext(filename)
# txt文件名,与对应图片名只有后缀名不一样
txt_name = head + ".txt"
f_txt = open(os.path.join(arg.save_dir, txt_name), 'w')
for ann in data['annotations']:
if ann['image_id'] == img_id:
box = convert((img_width, img_height), ann["bbox"])
# 写入txt,共5个字段
f_txt.write("%s %s %s %s %s\n" % (
id_map[ann["category_id"]], box[0], box[1], box[2], box[3]))
f_txt.close()