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 模型(注意:第一次使用需要点击模型管理,然后选择模型,下载对应的模型权重)。如果以上那个目录有同名文件(微调的模型)的话就会自动加载,然后就可以利用微调后的模型标注了。
看到启动 ISAT 的命令行输出 - load fine_tuning model ! 说明成功加载微调模型。
参考了下面的内容
Train/Fine-Tune Segment Anything 2 (SAM 2) in 60 Lines of Code