基于Segment Anything(SAM)的AutoMaskCOCO图像轮廓语义分割标注数据集

探索AutoMaskCOCO，这款独一无二的图像分割工具，将Segment Anything、COCO Viewer和SALT融合成一体。它实现了一站式图像标注与COCO数据集生成，让数据准备变得前所未有地快速、高效，为计算机视觉研究与开发注入强大动力！

Youcheng_Zong

3584人浏览 · 2023-04-14 19:30:45

Youcheng_Zong · 2023-04-14 19:30:45 发布

@[TOC](基于Segment Anything(SAM)的AutoMaskCOCO图像轮廓语义分割标注数据集)

AutoMaskCOCO

AutoMaskCOCO是一个简单易用的图像分割数据集标注工具，它结合了Segment Anything、COCO Viewer和Segment Anything Labelling Tool (SALT) 三个开源项目的优势，为用户提供了一个自动生成掩码并创建COCO数据集的一站式解决方案。通过这个项目，我们旨在简化图像分割数据集标注的过程，为计算机视觉研究人员和开发者节省宝贵的时间。

Segment Anything为AutoMaskCOCO提供了灵活的分割与掩码。
COCO Viewer项目为AutoMaskCOCO带来了一个直观的界面，方便用户查看和管理COCO数据集，对图像和标注进行有效的检查。
Segment Anything Labelling Tool (SALT) 为AutoMaskCOCO提供了一个用户友好的标注界面，让用户能够轻松地对图像进行分割、标注和编辑。

AutoMaskCOCO项目利用标点自动生成掩码，降低了图像分割数据标注的难度，同时还可以自动输出COCO数据集格式的结果，为计算机视觉任务的数据准备工作提供便利。

示例

How it Works Gif! Panel Jpg!
COCOViewerPanel Jpg!

准备工作

克隆相关项目：

git clone https://github.com/917940234/SAM-Tool.git
git clone https://github.com/facebookresearch/segment-anything.git

下载预训练的SAM模型：sam_vit_h_4b8939.pth 并将其放入segment-anything项目主目录。
准备数据集：将图像放置在<dataset_path>/images/中，并在<dataset_path>下创建名为embeddings的空文件夹。
将SAM-Tool项目中的helpers文件夹复制到segment-anything项目的主目录。
安装环境：

conda install pytorch==1.12.1 torchvision==0.13.1 torchaudio==0.12.1 cudatoolkit=11.6 -c pytorch -c conda-forge

提取图像Embeddings

运行extract_embeddings.py脚本以提取图像的embeddings：

python helpers\extract_embeddings.py --checkpoint-path sam_vit_h_4b8939.pth --dataset-folder <dataset_path>

参数说明：

checkpoint-path：预训练SAM模型的路径。
dataset-folder：数据集路径。
device：默认为cuda，如无GPU，可使用cpu（速度较慢）。
运行完成后，<dataset_path>/embeddings文件夹下会生成相应的npy文件。

转换模型格式

使用generate_onnx.py脚本将.pth文件转换为.onnx模型文件：

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模型供后续使用。

标注图像

将生成的sam_onnx.onnx模型复制到SAM-Tool项目主目录，然后运行segment_anything_annotator.py进行标注：

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：数据集的类别（每个类别以,分割，不要有空格）。
在对象位置点击鼠标左键添加掩码，点击右键去掉该位置掩码。最后生成的标注文件为COCO格式，保存在<dataset_path>/annotations.json。

查看标注结果

在SAM-Tool项目中运行cocoviewer.py查看全部的标注结果：

python cocoviewer.py -i <dataset_path> -a <dataset_path>\annotations.json

自定义设置

修改标注框线条宽度：在salt/display_utils.py中，找到class DisplayUtils，修改self.box_width的值。
修改标注文本格式：在salt/display_utils.py中，找到def draw_box_on_image函数，修改text、txt_color、font等参数。
标注快捷键：