利用SAM实现自动标注

利用SAM实现自动标注

一、下载安装

1.1 下载SAM-Tool和SAM

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

1.2 下载SAM模型文件

wget https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth

1.3 安装SAM

进入SAM文件夹：cd segment-anything
安装SAM：pip install -e .

二、配置项目

（1）将SAM工具中的helper文件夹移动到SAM主文件夹中：cp -r helpers/ …/segment-anything/
（2）构建数据文件夹格式为dataset_path/images和dataset_path/embeddings(images里用来放图片)

三、提取信息

（1）提取图片embedding：python helpers/extract_embeddings.py --checkpoint-path sam_vit_h_4b8939.pth --dataset-folder 你的数据集地址 --device cuda
（2）运行完在embeddings下会有对应的npy文件

四、获得SAM onnx文件

（1）首先环境中要有onnx和onnxruntime(-gpu)
（2）检测图像的文件(可调整后面的图片高/宽)：python helpers/generate_onnx.py --checkpoint-path sam_vit_h_4b8939.pth --onnx-model-path ./sam_onnx.onnx --orig-im-size 720 1280
（3）运行完会有对应的sam_onnx.onnx文件，将其移到SAM工具主文件夹中：cp sam_onnx.onnx …/SAM-Tool/

五、标注

（1）opencv库要删掉opencv-python，改用headless版的：pip install opencv-python-headless
（2）调用标注程序命令为(后面可改标注类别，用逗号隔开)：python segment_anything_annotator.py --onnx-model-path sam_onnx.onnx --dataset-path …\segment-anything\dataset --categories cement,metal,plastics
（3）标注框出现后开始标注，首先点击目标，然后添加对象，最后选定类别，接下来下一张，标完所有就保存，最后会在数据文件夹下出现annotations.json的coco格式文件
（4）查看标注文件：python cocoviewer.py -i …/segment-anything/dataset/ -a …/segment-anything/dataset/annotations.json