SAM2分割模型微调指南

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

你可能会问，如果 SAM 可以分割任何东西，为什么我们还需要重新训练它？答案是 SAM 在常见物体方面非常出色，但在罕见或特定领域的任务上表现不佳。

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

本教程演示了如何在仅 60 行代码（不包括标注和导入）中对新数据进行微调 SAM2。

完整的训练脚本可以在这里找到。

NSDT工具推荐： Three.js AI纹理开发包 - YOLO合成数据生成器 - GLTF/GLB在线编辑 - 3D模型格式在线转换 - 可编程3D场景编辑器 - REVIT导出3D模型插件 - 3D模型语义搜索引擎 - Three.js虚拟轴心开发包 - 3D模型在线减面 - STL模型在线切割

1、SAM的工作原理

SAM 的主要工作方式是获取图像和图像中的一个点，并预测包含该点的片段的掩码。这种方法无需人工干预即可实现完整的图像分割，并且对片段的类别或类型没有任何限制（如上一篇文章所述）。

使用 SAM 进行完整图像分割的过程：

选择图像中的一组点
使用 SAM 预测包含每个点的片段
将生成的片段组合成单个地图

虽然 SAM 还可以利用其他输入，例如掩码或边界框，但这些输入主要与涉及人工输入的交互式分割相关。在本教程中，我们将重点介绍全自动分割，并且仅考虑单点输入。

有关该模型的更多详细信息，请访问SAM项目网站。

2、下载 SAM2 并设置环境

SAM2 可从这里下载。如果你不想复制训练代码，也可以下载我分叉的版本，其中已经包含TRAIN.py脚本。

然后按照 github 存储库上的安装说明进行操作。通常，你需要 Python >=3.11 和 PyTorch。

此外，我们将使用 OpenCV，可以使用以下方式安装：

pip install opencv-python

你还需要从这里下载预训练模型，你可以从多种模型中进行选择，所有模型都与本教程兼容。我建议使用训练速度最快的小型模型。

下一步是下载用于微调模型的数据集。在本教程中，我们将使用 LabPics1 数据集来分割材料和液体。你可以从这里下载数据集。

4、准备数据读取器

我们需要编写的第一件事是数据读取器。它将读取并准备网络数据。

数据读取器需要生成：

图像
图像中所有片段的蒙版。
以及每个蒙版内的随机点

让我们从加载依赖项开始：

import numpy as np
import torch
import cv2
import os
from sam2.build_sam import build_sam2
from sam2.sam2_image_predictor import SAM2ImagePredictor

接下来我们列出数据集中的所有图像：

data_dir=r"LabPicsV1//" # Path to LabPics1 dataset folder
data=[] # list of files in dataset
for ff, name in enumerate(os.listdir(data_dir+"Simple/Train/Image/")):  # go over all folder annotation
    data.append({"image":data_dir+"Simple/Train/Image/"+name,"annotation":data_dir+"Simple/Train/Instance/"+name[:-4]+".png"})

现在介绍加载训练批次的主要函数。训练批次包括：一张随机图像、属于该图像的所有分割掩码以及每个掩码中的一个随机点：

def read_batch(data): # read random image and its annotaion from  the dataset (LabPics)

   #  select image

        ent  = data[np.random.randint(len(data))] # choose random entry
        Img = cv2.imread(ent["image"])[...,::-1]  # read image
        ann_map = cv2.imread(ent["annotation"]) # read annotation

   # resize image

        r = np.min([1024 / Img.shape[1], 1024 / Img.shape[0]]) # scalling factor
        Img = cv2.resize(Img, (int(Img.shape[1] * r), int(Img.shape[0] * r)))
        ann_map = cv2.resize(ann_map, (int(ann_map.shape[1] * r), int(ann_map.shape[0] * r)),interpolation=cv2.INTER_NEAREST)

   # merge vessels and materials annotations

        mat_map = ann_map[:,:,0] # material annotation map
        ves_map = ann_map[:,:,2] # vessel  annotaion map
        mat_map[mat_map==0] = ves_map[mat_map==0]*(mat_map.max()+1) # merged map

   # Get binary masks and points

        inds = np.unique(mat_map)[1:] # load all indices
        points= []
        masks = [] 
        for ind in inds:
            mask=(mat_map == ind).astype(np.uint8) # make binary mask
            masks.append(mask)
            coords = np.argwhere(mask > 0) # get all coordinates in mask
            yx = np.array(coords[np.random.randint(len(coords))]) # choose random point/coordinate
            points.append([[yx[1], yx[0]]])
        return Img,np.array(masks),np.array(points), np.ones([len(masks),1])

该函数的第一部分是选择一个随机图像并加载它：

ent  = data[np.random.randint(len(data))] # choose random entry
Img = cv2.imread(ent["image"])[...,::-1]  # read image
ann_map = cv2.imread(ent["annotation"]) # read annotation
Note that OpenCV reads images as BGR while SAM expects images as RGB, using […,::-1] to change the image from BGR to RGB.

请注意，OpenCV 将图像读取为 BGR，而 SAM 需要 RGB 图像。通过使用 […,::-1]，我们将图像从 BGR 更改为 RGB。

SAM 预计图像大小不超过 1024，因此我们将调整图像和注释图的大小至此大小：

r = np.min([1024 / Img.shape[1], 1024 / Img.shape[0]]) # scalling factor
Img = cv2.resize(Img, (int(Img.shape[1] * r), int(Img.shape[0] * r)))
ann_map = cv2.resize(ann_map, (int(ann_map.shape[1] * r), int(ann_map.shape[0] * r)),interpolation=cv2.INTER_NEAREST)

这里的一个重点是，在调整标注图 (ann_map) 的大小时，我们使用 INTER_NEAREST 模式（最近邻）。在注释图中，每个像素值都是其所属段的索引。因此，使用不会向地图引入新值的调整大小方法非常重要。

下一个块特定于 LabPics1 数据集的格式。标注图 (ann_map) 在一个通道中包含图像中血管的分割图，在另一个通道中包含材料注释的另一个图。我们将把它们合并成一张地图。

  mat_map = ann_map[:,:,0] # material annotation map
  ves_map = ann_map[:,:,2] # vessel  annotaion map
  mat_map[mat_map==0] = ves_map[mat_map==0]*(mat_map.max()+1) # merged map

这为我们提供了一个映射 (mat_map)，其中每个像素的值是其所属段的索引（例如：所有值为 3 的单元格都属于段 3）。我们希望将其转换为一组二进制掩码 (0/1)，其中每个掩码对应不同的段。此外，我们希望从每个掩码中提取一个点。

inds = np.unique(mat_map)[1:] # list of all indices in map
points= [] # list of all points (one for each mask)
masks = [] # list of all masks
for ind in inds:
            mask = (mat_map == ind).astype(np.uint8) # make binary mask for index ind
            masks.append(mask)
            coords = np.argwhere(mask > 0) # get all coordinates in mask
            yx = np.array(coords[np.random.randint(len(coords))]) # choose random point/coordinate
            points.append([[yx[1], yx[0]]])
return Img,np.array(masks),np.array(points), np.ones([len(masks),1])

就是这样！我们得到了图像（Img）中与图像中的片段（masks）相对应的二进制掩码列表，以及每个掩码中单个点的坐标（points）。

一批训练数据的示例：1）图像。2）片段掩码列表。3）每个掩码内的单个点（仅为可视化标记为红色）

5、加载 SAM 模型

现在让我们加载网络：

sam2_checkpoint = "sam2_hiera_small.pt" # path to model weight
model_cfg = "sam2_hiera_s.yaml" # model config
sam2_model = build_sam2(model_cfg, sam2_checkpoint, device="cuda") # load model
predictor = SAM2ImagePredictor(sam2_model) # load net

首先，我们在 sam2_checkpoint 参数中设置模型权重的路径。我们之前从这里下载了权重。“sam2_hiera_small.pt”指的是小模型，但代码适用于你选择的任何模型。无论选择哪种模型，都需要在 model_cfg 参数中设置相应的配置文件。配置文件已经位于主存储库的子文件夹“sam2_configs/”中。

6、SAM模型的一般结构

在设置训练参数之前，我们需要了解 SAM 模型的基本结构。

SAM 由三部分组成：1)图像编码器，2) 提示编码器，3) 掩码解码器。

图像编码器负责处理图像并创建表示图像的嵌入。这部分由 VIT 转换器组成，是网络的最大组件。我们通常不想训练它，因为它已经提供了良好的表示，训练将需要大量资源。

提示编码器处理网络的额外输入，在我们的例子中是输入点。

掩码解码器获取图像编码器和提示编码器的输出并生成最终的分割掩码。一般来说，我们只想训练掩码解码器和提示编码器。这些部分很轻量，可以用适中的 GPU 快速微调。

7、设置训练参数

我们可以通过设置来启用掩码解码器和提示编码器的训练：

predictor.model.sam_mask_decoder.train(True) # enable training of mask decoder 
predictor.model.sam_prompt_encoder.train(True) # enable training of prompt encoder

接下来，我们定义标准的adamW优化器：

optimizer=torch.optim.AdamW(params=predictor.model.parameters(),lr=1e-5,weight_decay=4e-5)

我们还将使用混合精度训练，这是一种更节省内存的训练策略：

scaler = torch.cuda.amp.GradScaler() # set mixed precision

8、主训练循环

现在让我们构建主训练循环。第一部分是读取和准备数据：

for itr in range(100000):
    with torch.cuda.amp.autocast(): # cast to mix precision
            image,mask,input_point, input_label = read_batch(data) # load data batch
            if mask.shape[0]==0: continue # ignore empty batches
            predictor.set_image(image) # apply SAM image encoder to the image

首先，我们将数据转换为混合精度以实现高效训练：

with torch.cuda.amp.autocast():

接下来，我们使用之前创建的读取器函数来读取训练数据：

image,mask,input_point, input_label = read_batch(data)

我们将加载的图像传递给图像编码器（网络的第一部分）：

predictor.set_image(image)

接下来，我们使用网络提示编码器处理输入点：

  mask_input, unnorm_coords, labels, unnorm_box = predictor._prep_prompts(input_point, input_label, box=None, mask_logits=None, normalize_coords=True)
  sparse_embeddings, dense_embeddings = predictor.model.sam_prompt_encoder(points=(unnorm_coords, labels),boxes=None,masks=None,)

请注意，在这一部分我们也可以输入框或掩码，但我们不会使用这些选项。

现在我们对提示（点）和图像都进行了编码，我们最终可以预测分割掩码：

batched_mode = unnorm_coords.shape[0] > 1 # multi mask prediction
high_res_features = [feat_level[-1].unsqueeze(0) for feat_level in predictor._features["high_res_feats"]]
low_res_masks, prd_scores, _, _ = predictor.model.sam_mask_decoder(image_embeddings=predictor._features["image_embed"][-1].unsqueeze(0),image_pe=predictor.model.sam_prompt_encoder.get_dense_pe(),sparse_prompt_embeddings=sparse_embeddings,dense_prompt_embeddings=dense_embeddings,multimask_output=True,repeat_image=batched_mode,high_res_features=high_res_features,)
prd_masks = predictor._transforms.postprocess_masks(low_res_masks, predictor._orig_hw[-1])# Upscale the masks to the original image resolution

此代码的主要部分是 model.sam_mask_decoder，它运行网络的 mask_decoder 部分并生成分割掩码 (low_res_masks) 及其分数 (prd_scores)。

这些掩码的分辨率低于原始输入图像，并在 postprocess_masks 函数中调整为原始输入大小。

这为我们提供了网络的最终预测：我们使用的每个输入点的 3 个分割掩码 (prd_masks) 和掩码分数 (prd_scores)。prd_masks 包含每个输入点的 3 个预测掩码，但我们只会使用每个点的第一个掩码。prd_scores 包含网络认为每个掩码有多好的分数（或它在预测中的确定性）。

9、损失函数

现在我们有了净预测，我们可以计算损失了。首先，我们计算分割损失，这意味着预测的掩码与地面真实掩码相比有多好。为此，我们使用标准交叉熵损失。

首先，我们需要使用 sigmoid 函数将预测掩码 (prd_mask) 从 logits 转换为概率：

prd_mask = torch.sigmoid(prd_masks[:, 0])# Turn logit map to probability map

接下来我们将真实情况掩码转换为 torch 张量：

prd_mask = torch.sigmoid(prd_masks[:, 0])# Turn logit map to probability map

最后，我们使用基本事实（gt_mask）和预测概率图（prd_mask）手动计算交叉熵损失（seg_loss）：

seg_loss = (-gt_mask * torch.log(prd_mask + 0.00001) - (1 - gt_mask) * torch.log((1 - prd_mask) + 0.00001)).mean() # cross entropy loss

（我们添加 0.0001 以防止对数函数因零值而爆炸）。

分数损失（可选）

除了掩码之外，网络还预测每个预测掩码的好坏分数。训练这部分不太重要，但很有用。要训练这部分，我们首先需要知道每个预测掩码的真实分数是多少。也就是说，预测的掩码实际上有多好。我们将通过使用交并比 (IOU) 指标比较 GT 掩码和相应的预测掩码来实现这一点。IOU 只是两个掩码之间的重叠，除以两个掩码的总面积。首先，我们计算预测和 GT 掩码之间的交集（它们重叠的区域）：

inter = (gt_mask * (prd_mask > 0.5)).sum(1).sum(1)

我们使用阈值 (prd_mask > 0.5) 将预测掩码从概率转换为二元掩码。

接下来，我们通过将交集除以预测掩码和 gt 掩码的组合面积（并集）来获得 IOU：

iou = inter / (gt_mask.sum(1).sum(1) + (prd_mask > 0.5).sum(1).sum(1) - inter)

我们将使用 IOU 作为每个 mask 的真实分数，并将预测分数与我们刚刚计算的 IOU 之间的绝对差作为分数损失。

score_loss = torch.abs(prd_scores[:, 0] - iou).mean()

最后，我们合并分割损失和分数损失（给予第一个更高的权重）：

loss = seg_loss+score_loss*0.05  # mix losses

10、反向传播和保存模型

一旦我们得到损失，一切都完全标准化了。我们使用之前制作的优化器计算反向传播并更新权重：

predictor.model.zero_grad() # empty gradient
scaler.scale(loss).backward()  # Backpropogate
scaler.step(optimizer)
scaler.update() # Mix precision

我们还希望每 1000 步保存一次训练好的模型：

if itr%1000==0: torch.save(predictor.model.state_dict(), "model.torch") # save model

由于我们已经计算了 IOU，我们可以将其显示为移动平均值，以查看模型预测随时间的改善程度：


if itr==0: mean_iou=0
mean_iou = mean_iou * 0.99 + 0.01 * np.mean(iou.cpu().detach().numpy())
print("step)",itr, "Accuracy(IOU)=",mean_iou)

就这样，我们用不到 60 行代码（不包括注释和导入）训练/微调了 Segment-Anything 2。经过大约 25,000 步后，你应该会看到重大改进。

该模型将保存到“model.torch”。可以在这里找到完整的训练代码。

11、推理：加载和使用经过训练的模型：

现在模型已经过微调，让我们用它来分割图像。

我们将使用以下步骤执行此操作：

加载我们刚刚训练的模型。
给模型一张图像和一堆随机点。对于每个点，网络将预测包含该点和分数的分割掩码。
获取这些掩码并将它们拼接成一个分割图。

执行此操作的完整代码可在这里找到。

首先，我们加载依赖项并将权重转换为 float16，这使得模型运行速度更快（仅适用于推理）。

# use bfloat16 for the entire script (memory efficient)
torch.autocast(device_type="cuda", dtype=torch.bfloat16).__enter__()

接下来，我们加载示例图像和我们想要分割的图像区域的蒙版（下载图像/蒙版）：

image_path = r"sample_image.jpg" # path to image
mask_path = r"sample_mask.png" # path to mask, the mask will define the image region to segment
def read_image(image_path, mask_path): # read and resize image and mask
        img = cv2.imread(image_path)[...,::-1]  # read image as rgb
        mask = cv2.imread(mask_path,0) # mask of the region we want to segment
        
        # Resize image to maximum size of 1024

        r = np.min([1024 / img.shape[1], 1024 / img.shape[0]])
        img = cv2.resize(img, (int(img.shape[1] * r), int(img.shape[0] * r)))
        mask = cv2.resize(mask, (int(mask.shape[1] * r), int(mask.shape[0] * r)),interpolation=cv2.INTER_NEAREST)
        return img, mask
image,mask = read_image(image_path, mask_path)

在我们要分割的区域内随机取样30个点：

num_samples = 30 # number of points/segment to sample
def get_points(mask,num_points): # Sample points inside the input mask
        points=[]
        for i in range(num_points):
            coords = np.argwhere(mask > 0)
            yx = np.array(coords[np.random.randint(len(coords))])
            points.append([[yx[1], yx[0]]])
        return np.array(points)
input_points = get_points(mask,num_samples)

首先，加载标准SAM模型（与训练中相同）

# Load model you need to have pretrained model already made
sam2_checkpoint = "sam2_hiera_small.pt" 
model_cfg = "sam2_hiera_s.yaml" 
sam2_model = build_sam2(model_cfg, sam2_checkpoint, device="cuda")
predictor = SAM2ImagePredictor(sam2_model)

接下来，加载我们刚刚训练的模型的权重（model.torch）：

predictor.model.load_state_dict(torch.load("model.torch"))

运行微调模型来预测我们选择的每个点的掩码：

with torch.no_grad(): # prevent the net from caclulate gradient (more efficient inference)
        predictor.set_image(image) # image encoder
        masks, scores, logits = predictor.predict(  # prompt encoder + mask decoder
            point_coords=input_points,
            point_labels=np.ones([input_points.shape[0],1])
        )

现在我们有了预测掩码及其分数的列表。我们希望以某种方式将它们拼接成一个一致的分割图。但是，许多掩码重叠并且可能彼此不一致。由于我们随机选择了这些点，因此很可能有些点会落在同一段中。

拼接方法很简单，我们将根据预测的分数对预测的掩码进行排序：

np_masks = np.array(masks[:,0].cpu().numpy()) # convert from torch to numpy
np_scores = scores[:,0].float().cpu().numpy() # convert from torch to numpy
shorted_masks = np_masks[np.argsort(np_scores)][::-1] # arrange mask according to score

现在让我们创建一个空的分割图和占用图：

seg_map = np.zeros_like(shorted_masks[0],dtype=np.uint8)
occupancy_mask = np.zeros_like(shorted_masks[0],dtype=bool)

接下来，我们将蒙版逐一（从高分到低分）添加到分割图中。我们只会在蒙版与之前添加的蒙版一致时才添加蒙版，也就是说，只有当我们想要添加的蒙版与已占用区域的重叠度小于 15% 时，我们才会添加蒙版。

for i in range(shorted_masks.shape[0]):
    mask = shorted_masks[i]
    if (mask*occupancy_mask).sum()/mask.sum()>0.15: continue 
    mask[occupancy_mask]=0
    seg_map[mask]=i+1
    occupancy_mask[mask]=1

就是这样。

seg_mask 现在包含预测的分割图，每个分割具有不同的值，背景为 0。

我们可以使用以下方法将其转换为颜色图：

rgb_image = np.zeros((seg_map.shape[0], seg_map.shape[1], 3), dtype=np.uint8)
for id_class in range(1,seg_map.max()+1):
    rgb_image[seg_map == id_class] = [np.random.randint(255), np.random.randint(255), np.random.randint(255)]

并显示：

cv2.imshow("annotation",rgb_image)
cv2.imshow("mix",(rgb_image/2+image/2).astype(np.uint8))
cv2.imshow("image",image)
cv2.waitKey()

使用微调 SAM2 的分割结果示例

完整推理代码可在此处获取。

原文链接：SAM2分割模型微调 - BimAnt