利用SAM半自动标注--微调SAM2

提亚斯的酒杯

已于 2024-09-09 11:03:29 修改

阅读量3.7k

点赞数 23

文章标签：人工智能深度学习机器学习

于 2024-09-08 22:22:02 首次发布

本文链接：https://blog.youkuaiyun.com/weixin_45586769/article/details/142029798

版权

SAM2（Segment Anything 2）是 Meta 推出的一款新模型，旨在对图像中的任何内容进行分割，而不局限于特定的类别或领域。该模型的独特之处在于其训练数据规模：1100 万张图像和 110 亿个掩码。这种广泛的训练使 SAM2 成为训练新图像分割任务的强大起点。

但是在利用 ISAT 打标签的过程中发现 SAM2 对裂缝的检测精度很差，这时候就需要微调SAM，让它学习标注。

即使在 SAM 给出的结果不足的情况下，仍然可以通过在新数据上对其进行微调来显着提高模型的能力。在许多情况下，这将比从头开始训练模型占用更少的训练数据并获得更好的结果。

1 配置 SAM 2 环境

在这里下载 segment-anything-2 代码，参照 Installation 部分，安装 SAM2。

2 labelme 标注一部分标签，将 labelme 格式标签(json)转换为png格式

png 格式的标签是灰度图，通过不同的灰度值，表示不同类别的目标。

例如下图，左图是原始图片，右图是对应的 png 标签，中间的图像是对 png 标签可视化。

图中绿色部分是裂缝，在标签(右图)中这部分的灰度值为1，

蓝色部分是条修，在标签(右图)中这部分的灰度值为2，

背景部分的灰度值当然为0了。

labelme 格式标签转换为 png 格式标签的代码可在这里找到。将label_map、input_dir、output_dir改成自己的。

# 定义类别和对应的灰度值
label_map = {
    "background": 0,
    "songsan": 1,
    "liefeng": 2,
    "junlie": 3,
    "kuaixiu": 4,
    "tiaoxiu": 5,
    "aokeng": 6
}
# 示例用法
input_dir = r'./ZH/png2json/'  # 输入目录
output_dir = r'./ZH/json2png/'  # 输出目录
process_all_json_files(input_dir, output_dir)

对png标签进行检查（非必要）

通过这个代码可以将png灰度图转换为彩色图片，用来检查标签是否准确。

需要更改下面的代码。

# 定义类别数
num_classes = 7

# 创建一个颜色映射表
color_map = np.array([
    [0, 0, 0],       # 背景 Background
    [255, 0, 0],     # 类别1 Class 1  红 songsan
    [0, 255, 0],     # 类别2 Class 2  绿 liefeng
    [0, 0, 255],     # 类别3 Class 3  蓝 junlie
    [255, 255, 0],   # 类别4 Class 4  黄 kuaixiu
    [0, 255, 255],   # 类别5 Class 5  青 tiaoxiu
    [255, 0, 255]    # 类别6 Class 6  洋红 aokeng
    # 类别7 Class 7 可以选择其他颜色，比如 [128, 128, 128]
])

# 文件夹路径
folder_path = r"./out_t"
out_folder_path = r'./'

3 微调 SAM 2

微调 SAM 2 的脚本，在这里。

需要预先在这里下载 SAM2 的权重，比如我下载了 sam2_hiera_small.pt

需要根据自己的情况更改下面的内容（从169行开始）。

sam2_checkpoint 与 model_cfg 要对应

class_id 指的是需要微调的类别。比如 SAM 对裂缝的识别效果很差，在我的标签中裂缝对应的灰度值为2，所以填一个2就行了。若想训练多个类别，填一个数组。

     # Read data
     data_dir=r'/home/user/works/LLL-4090/segment-anything-2-main/data/road-dataset/' # Path to dataset
     data=[] # list of files in dataset
     #for ff, name in enumerate(os.listdir(data_dir+"images/")):  # go over all folder annotation
     image_path = os.path.join(data_dir, 'images/')
     for ff, name in enumerate(os.listdir(image_path)):  # go over all folder annotation
          data.append({"image":data_dir+"images/"+name,"annotation":data_dir+"labels-png/"+name[:-4]+".png"})

     # 单独微调灰度为2的类别
     class_id = np.array([2])
     #class_id = np.array([1,2,3]) # 微调多个类别

     save_path = '/home/user/works/LLL-4090/segment-anything-2-main/runs'

     # Load model
     sam2_checkpoint = "/home/user/works/LLL-4090/segment-anything-2-main/data/sam2_hiera_small.pt" # path to model weight
     model_cfg = "sam2_hiera_s.yaml" # model config
     predictor = load_model(sam2_checkpoint, model_cfg, device='cuda:1')

下图是训练的结果，save model in 后面是权重的保存路径，best iou 是当前权重的 iou。

4 对标注软件 ISAT 进行修改，以加载微调后的模型

ISAT结合了Meta推出的Segment Anything Model（SAM），是一款功能强大且灵活的图像语义分割和打标签工具。在这里下载。

这篇博客详细介绍了ISAT，请参考安装: 分割——半自动打标签工具ISAT详细安装及使用教程

配置好环境后，修改 ISAT_with_segment_anythingmaster/ISAT/segment_any/segment_any.py

103行，添加以下内容，

        sam = sam.eval().to(self.model_dtype)

        sam.to(device=self.device)
        self.predictor_with_point_prompt = SamPredictor(sam)

        ''' 以下是添加的代码'''

        # 加载微调模型
        base_dir = os.path.dirname(checkpoint) 
        file_name = os.path.basename(checkpoint)
        fine_tuning = os.path.join(base_dir, 'fine_tuning', file_name)
        if os.path.exists(fine_tuning):
            self.predictor_with_point_prompt.model.load_state_dict(torch.load(fine_tuning, map_location=self.device))
            print('- load fine_tuning model !')
        
        '''以上是添加的代码'''

        print('- loaded')
        self.image = None

将上一步微调好的模型，改名为 sam2_hiera_small.pt ，也就是你微调的模型名称，然后复制到ISAT_with_segment_anything-master/ISAT/checkpoints/fine_tuning/sam2_hiera_small.pt

在启动 ISAT 后，正常加载 sam2_hiera_small 模型（注意：第一次使用需要点击模型管理，然后选择模型，下载对应的模型权重）。如果以上那个目录有同名文件（微调的模型）的话就会自动加载，然后就可以利用微调后的模型标注了。