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

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

AutoMaskCOCO

AutoMaskCOCO是一个简单易用的图像分割数据集标注工具,它结合了Segment AnythingCOCO ViewerSegment 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!

准备工作

  1. 克隆相关项目:
git clone https://github.com/917940234/SAM-Tool.git
git clone https://github.com/facebookresearch/segment-anything.git
  1. 下载预训练的SAM模型:sam_vit_h_4b8939.pth 并将其放入segment-anything项目主目录。
  2. 准备数据集:将图像放置在<dataset_path>/images/中,并在<dataset_path>下创建名为embeddings的空文件夹。
  3. SAM-Tool项目中的helpers文件夹复制到segment-anything项目的主目录。
  4. 安装环境:
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

自定义设置

  1. 修改标注框线条宽度: 在salt/display_utils.py中,找到class DisplayUtils,修改self.box_width的值。
  2. 修改标注文本格式: 在salt/display_utils.py中,找到def draw_box_on_image函数,修改texttxt_colorfont等参数。
  3. 标注快捷键:
  • Esc:退出app
  • a:前一张图片
  • d:后一张图片
  • a:前十张图片
  • d:后十张图片
  • k:调低透明度
  • l:调高透明度
  • n:添加对象
  • r:重置
  • Ctrl+s:保存之前的操作
  • Ctrl+z:撤销对象

环境

参考项目中environment.yamlrequirement.py

通过以上步骤,可以使用AutoMaskCOCO标注工具为图像数据集生成COCO格式的标注。如有需要,可以根据自己的需求调整设置以优化标注过程。

参考与引用:
Segment Anything
Segment Anything Labelling Tool (SALT)
COCO Viewer
使用Segment Anything(SAM)模型进行自动标注

评论 11
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值