使用segment-anything将目标检测label转换为语义分割label

一、segment-anything简介

segment-anything是facebookresearch团队开发的一套无须训练，根据提示对图像自动分割的工具，提示可以是一个点，也可以是一个矩形框。根据主页介绍，它是使用1100万图片和1.1亿mask训练而成，在各类图片分割中都有很强的性能。

研究团队还在此基础上开发出了Segment-Anything-Model 2，简称SAM2，SAM2不仅可以对图像进行分割，还可以对视频进行分割。

segment-anything项目Github地址为：https://github.com/facebookresearch/segment-anything，
在线演示地址为：https://segment-anything.com/demo，可以在该网站上传一张本地图片，进行试验：

SAM2项目Github地址为：https://github.com/facebookresearch/sam2。

二、segment-anything安装

2.1安装方法

segment-anything要求环境：python>=3.8，pytorch>=1.7，torchvision>=0.8。
安装方法有两种，一种为使用pip命令安装，安装命令如下：

pip install git+https://github.com/facebookresearch/segment-anything.git

一种为下载项目源代码进行本地安装，如下：

git clone git@github.com:facebookresearch/segment-anything.git
cd segment-anything; pip install -e .

2.2预训练模型下载

使用segment-anything需要下载对应的预训练模型，按照参数量从小到大分为：vit_b、vit_l、vit_h，大小分别为：360M、1.2G，2.5G。
三种模型对显卡显存的要求也依次增高，根据实际测试，8G显存的显卡可以加载运行vit_b和vit_l，无法加载vit_h模型。

模型下载官方地址如下：
vit_b：https://dl.fbaipublicfiles.com/segment_anything/sam_vit_b_01ec64.pth
vit_l：https://dl.fbaipublicfiles.com/segment_anything/sam_vit_l_0b3195.pth
vit_h：https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth。

CSDN免费下载地址如下：
vit_b：https://download.csdn.net/download/weixin_46846685/91510193
vit_l：https://download.csdn.net/download/weixin_46846685/91510428，https://download.csdn.net/download/weixin_46846685/91510431（CSDN上传文件限制，分为两部分）
vit_h：https://download.csdn.net/download/weixin_46846685/91510458，https://download.csdn.net/download/weixin_46846685/91553916，https://download.csdn.net/download/weixin_46846685/91553931（CSDN上传文件限制，分为三部分）

三、将目标检测label转换为语义分割label

目标检测任务中的label为目标的矩形框坐标，而语义分割任务中的label为目标的轮廓。如果需要将目标检测任务转换为语义分割任务，那么就需要对图片进行重新标注。

这一过程需要耗费大量时间精力，所以本文给出了使用segment-anything进行自动化转换的代码，用于提高工作效率，但需要注意的是，实际分割效果需要根据场景进行测试评估，检查效果是否满足预期。

3.1示例代码

以YOLO的标签格式为例，展示如何根据目标检测的box坐标，将目标轮廓从图片中分割出来生成对应的mask。
代码如下：

import torch
import numpy as np
import cv2
from segment_anything import sam_model_registry, SamPredictor, SamAutomaticMaskGenerator
from PIL import Image
import os
import sysdevice = "cuda" if torch.cuda.is_available() else "cpu"
sam_checkpoint = "sam_vit_l_0b3195.pth"
model_type = "vit_l"sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)
sam.to(device)imagepath = 'image.jpg'
labelpath = 'label.txt'
combine = np.zeros([256, 256])
if os.path.getsize(labelpath) > 0:with open(labelpath, 'r') as f:labels=f.readlines()boxes = []for labeldata in labels:parts = labeldata.strip().split()x_center, y_center, w, h = (float(parts[1]), float(parts[2]), float(parts[3]), float(parts[4]))gtx1 = int((x_center - w / 2) * 256)gty1 = int((y_center - h / 2) * 256)gtx2 = int((x_center + w / 2) * 256)gty2 = int((y_center + h / 2) * 256)boxes.append([gtx1, gty1, gtx2, gty2])image = cv2.imread(imagepath)image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)predictor = SamPredictor(sam)predictor.set_image(image)for box in boxes:masks, scores, logits = predictor.predict(point_coords=None,point_labels=None,box=np.array(box),multimask_output=True,)mask = masks[scores[0].argmax()]combine[np.where(mask==1)]=1mask_output_path = 'mask.png'
mask_img = Image.fromarray(combine.astype(np.uint8))
mask_img.save(mask_output_path)

3.2代码说明

代码整体逻辑为：

1. 首先判断label文件是否为空，如果为空说明图片不包含目标，则生成一个全0的mask；
2. 如果不为空，则读取label的box信息，以box为提示，调用segment-anything对图片进行分割，上述代码假设目标只有1类，所以将目标区域填充为1，此处可根据需要进行修改；
3. 最后生成的mask，背景区域灰度值为0，目标区域灰度值为1，所以直接打开mask图片无法看到分割出的目标，如需可视化需要另做一些简单的图像处理。

调用segment-anything的注意事项：

1. segment-anything不支持直接输入box的列表，所以上述代码做了一个循环，循环输入box；
2. predictor.predict_torch方法据说可以接受多个box，但是经过实际测试，发现分割出的目标位置有所偏离，目前不知是什么原因。