从零到一：如何用SAM2-UNet构建你的首个医学图像分割项目

医学图像分割是计算机视觉在医疗领域的重要应用场景，它能够帮助医生更准确地识别和测量生物组织结构，为疾病诊断提供有力支持。近年来，随着深度学习技术的发展，基于U-Net架构的分割模型已经成为医学图像分析的标准工具。而Segment Anything Model 2（SAM2）的出现，为这一领域带来了新的突破。本文将详细介绍如何从零开始构建一个基于SAM2-UNet的医学图像分割项目，以息肉分割为例，展示完整的实现流程。

1. 环境配置与准备工作

在开始项目之前，我们需要搭建合适的开发环境。医学图像分割通常需要较强的计算能力，建议使用配备NVIDIA显卡的工作站或云服务器。以下是环境配置的具体步骤：

首先安装必要的Python包，建议使用Python 3.10或更高版本：

conda create -n sam2 python=3.10
conda activate sam2
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
pip install opencv-python matplotlib numpy tqdm

接下来安装SAM2-UNet的官方实现：

git clone https://github.com/facebookresearch/segment-anything-2.git
cd segment-anything-2
pip install -e .

对于医学图像处理，我们还需要安装一些专门的库：

pip install SimpleITK pydicom nibabel

环境配置完成后，下载预训练模型权重。SAM2提供了多种规模的模型，对于医学图像分割任务，推荐使用Hiera-Large版本：

wget https://dl.fbaipublicfiles.com/segment_anything_2/072824/sam2_hiera_large.pt -P checkpoints/

2. 数据集准备与预处理

医学图像数据通常以DICOM或NIfTI格式存储，我们需要将其转换为适合深度学习模型处理的格式。以息肉分割常用的Kvasir-SEG数据集为例，介绍数据预处理流程。

首先创建数据集目录结构：

data/
├── images/
├── masks/
├── train.txt
└── val.txt

使用Python脚本将原始数据转换为统一格式：

import os
import cv2
import pydicom
from glob import glob

def convert_dicom_to_png(dicom_path, output_path):
    ds = pydicom.dcmread(dicom_path)
    img = ds.pixel_array
    cv2.imwrite(output_path, img)

def prepare_kvasir_dataset(raw_dir, output_dir):
    os.makedirs(os.path.join(output_dir, 'images'), exist_ok=True)
    os.makedirs(os.path.join(output_dir, 'masks'), exist_ok=True)
    
    image_files = glob(os.path.join(raw_dir, '*.jpg'))
    for img_path in image_files:
        base_name = os.path.basename(img_path)
        mask_path = img_path.replace('.jpg', '_mask.jpg')
        
        # 处理图像
        img = cv2.imread(img_path)
        cv2.imwrite(os.path.join(output_dir, 'images', base_name), img)
        
        # 处理掩码
        mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)
        _, binary_mask = cv2.threshold(mask, 127, 255, cv2.THRESH_BINARY)
        cv2.imwrite(os.path.join(output_dir, 'masks', base_name), binary_mask)

医学图像数据通常需要特殊的增强处理，以下是一个增强流程的示例：

import albumentations as A

def get_augmentations():
    return A.Compose([
        A.RandomRotate90(),
        A.Flip(),
        A.ElasticTransform(alpha=120, sigma=120*0.05, alpha_affine=120*0.03),
        A.GridDistortion(),
        A.RandomBrightnessContrast(p=0.5),
        A.GaussianBlur(blur_limit=3, p=0.5),
    ], additional_targets={'mask': 'mask'})

3. SAM2-UNet模型架构解析

SAM2-UNet的核心创新在于将SAM2的Hiera骨干网络作为编码器，结合经典的U-Net解码器架构。这种设计既保留了SAM2强大的特征提取能力，又利用了U-Net在医学图像分割中的成熟优势。

模型的主要组件包括：

1. 编码器：基于SAM2预训练的Hiera骨干网络，采用层次结构捕获多尺度特征
2. 接收域块(RFBs)：减少通道数并增强轻量级特征
3. 适配器：实现参数高效的微调
4. 解码器：经典的U-Net设计，包含三个解码块

以下是模型构建的关键代码：

import torch
import torch.nn as nn
from sam2.build_sam import build_sam2

class SAM2UNet(nn.Module):
    def __init__(self, model_cfg='sam2_hiera_l.yaml', checkpoint='checkpoints/sam2_hiera_large.pt'):
        super().__init__()
        # 加载预训练的SAM2作为编码器
        self.encoder = build_sam2(model_cfg, checkpoint).image_encoder
        
        # 接收域块
        self.rfbs = nn.ModuleList([
            nn.Sequential(
                nn.Conv2d(144, 64, 1),
                nn.BatchNorm2d(64),
                nn.ReLU()
            ),
            nn.Sequential(
                nn.Conv2d(288, 64, 1),
                nn.BatchNorm2d(64),
                nn.ReLU()
            ),
            nn.Sequential(
                nn.Conv2d(576, 64, 1),
                nn.BatchNorm2d(64),
                nn.ReLU()
            ),
            nn.Sequential(
                nn.Conv2d(1152, 64, 1),
                nn.BatchNorm2d(64),
                nn.ReLU()
            )
        ])
        
        # 解码器
        self.decoder = nn.ModuleList([
            DecoderBlock(64*4, 256),
            DecoderBlock(256, 128),
            DecoderBlock(128, 64)
        ])
        
        self.seg_head = nn.Conv2d(64, 1, 1)
        
    def forward(self, x):
        # 编码器提取多尺度特征
        features = self.encoder(x)
        
        # 通过RFBs处理特征
        rf_features = []
        for f, rfb in zip(features, self.rfbs):
            rf_features.append(rfb(f))
        
        # 解码过程
        x = torch.cat(rf_features, dim=1)
        for decoder in self.decoder:
            x = decoder(x)
        
        return self.seg_head(x)

class DecoderBlock(nn.Module):
    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.block = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, 3, padding=1),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(),
            nn.Conv2d(out_channels, out_channels, 3, padding=1),
            nn.BatchNorm2d(out_channels),
            nn.ReLU()
        )
        self.upsample = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
        
    def forward(self, x):
        x = self.upsample(x)
        return self.block(x)

4. 模型训练与优化

训练医学图像分割模型需要考虑数据不平衡、小样本学习等特殊问题。以下是训练流程的关键步骤：

首先定义损失函数，医学图像分割常用Dice损失和交叉熵损失的组合：

class DiceBCELoss(nn.Module):
    def __init__(self, smooth=1.0):
        super().__init__()
        self.smooth = smooth
        
    def forward(self, pred, target):
        pred = torch.sigmoid(pred)
        intersection = (pred * target).sum()
        dice_loss = 1 - (2. * intersection + self.smooth) / 
                   (pred.sum() + target.sum() + self.smooth)
        bce = F.binary_cross_entropy_with_logits(pred, target)
        return dice_loss + bce

训练循环的实现：

def train_model(model, train_loader, val_loader, epochs=100, lr=1e-4):
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model = model.to(device)
    
    optimizer = torch.optim.AdamW(model.parameters(), lr=lr)
    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, 'max', patience=5)
    criterion = DiceBCELoss()
    
    best_score = 0
    for epoch in range(epochs):
        model.train()
        train_loss = 0
        for images, masks in tqdm(train_loader):
            images, masks = images.to(device), masks.to(device)
            
            optimizer.zero_grad()
            outputs = model(images)
            loss = criterion(outputs, masks)
            loss.backward()
            optimizer.step()
            
            train_loss += loss.item()
        
        # 验证阶段
        val_score = evaluate(model, val_loader, device)
        scheduler.step(val_score)
        
        print(f'Epoch {epoch+1}/{epochs} | Train Loss: {train_loss/len(train_loader):.4f} | Val Dice: {val_score:.4f}')
        
        # 保存最佳模型
        if val_score > best_score:
            best_score = val_score
            torch.save(model.state_dict(), 'best_model.pth')

评估指标的计算（以Dice系数为例）：

def dice_coeff(pred, target):
    smooth = 1.0
    pred = (pred > 0.5).float()
    intersection = (pred * target).sum()
    return (2. * intersection + smooth) / (pred.sum() + target.sum() + smooth)

def evaluate(model, loader, device):
    model.eval()
    total_dice = 0
    with torch.no_grad():
        for images, masks in loader:
            images, masks = images.to(device), masks.to(device)
            outputs = model(images)
            total_dice += dice_coeff(outputs, masks).item()
    return total_dice / len(loader)

5. 结果分析与实际应用

训练完成后，我们需要对模型性能进行全面评估。医学图像分割常用的评估指标包括：

指标名称	计算公式	意义
Dice系数	$\frac{2	X \cap Y
IoU	$\frac{	X \cap Y
敏感度	$\frac{TP}{TP + FN}$	识别阳性样本的能力
特异性	$\frac{TN}{TN + FP}$	识别阴性样本的能力

在实际应用中，我们可以使用训练好的模型进行息肉分割：

def predict_single_image(model, image_path, device='cuda'):
    # 加载图像
    image = cv2.imread(image_path)
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    original_size = image.shape[:2]
    
    # 预处理
    transform = A.Compose([
        A.Resize(256, 256),
        A.Normalize()
    ])
    transformed = transform(image=image)
    img_tensor = torch.from_numpy(transformed['image']).permute(2,0,1).unsqueeze(0).float().to(device)
    
    # 预测
    model.eval()
    with torch.no_grad():
        pred = model(img_tensor)
        pred = torch.sigmoid(pred)
    
    # 后处理
    pred_mask = (pred > 0.5).float().squeeze().cpu().numpy()
    pred_mask = cv2.resize(pred_mask, (original_size[1], original_size[0]))
    
    return pred_mask

对于医疗应用场景，我们还需要考虑模型的可解释性。可以使用Grad-CAM等方法可视化模型关注的重点区域：

class GradCAM:
    def __init__(self, model, target_layer):
        self.model = model
        self.target_layer = target_layer
        self.gradients = None
        self.activations = None
        
        target_layer.register_forward_hook(self.save_activations)
        target_layer.register_backward_hook(self.save_gradients)
    
    def save_activations(self, module, input, output):
        self.activations = output
    
    def save_gradients(self, module, grad_input, grad_output):
        self.gradients = grad_output[0]
    
    def __call__(self, x, target_class=None):
        self.model.zero_grad()
        
        # 前向传播
        output = self.model(x)
        if target_class is None:
            target_class = torch.argmax(output)
        
        # 反向传播
        one_hot = torch.zeros_like(output)
        one_hot[0][target_class] = 1
        output.backward(gradient=one_hot)
        
        # 计算权重
        weights = torch.mean(self.gradients, dim=(2,3), keepdim=True)
        cam = torch.sum(weights * self.activations, dim=1, keepdim=True)
        cam = F.relu(cam)
        cam = F.interpolate(cam, x.shape[2:], mode='bilinear', align_corners=False)
        
        # 归一化
        cam = cam - cam.min()
        cam = cam / cam.max()
        return cam.squeeze().cpu().numpy()

在实际部署时，可以考虑使用ONNX格式优化推理速度：

def export_to_onnx(model, sample_input, output_path='model.onnx'):
    torch.onnx.export(
        model,
        sample_input,
        output_path,
        input_names=['input'],
        output_names=['output'],
        dynamic_axes={
            'input': {0: 'batch_size', 2: 'height', 3: 'width'},
            'output': {0: 'batch_size', 2: 'height', 3: 'width'}
        },
        opset_version=11
    )
    
    # 验证ONNX模型
    import onnx
    onnx_model = onnx.load(output_path)
    onnx.checker.check_model(onnx_model)

6. 性能优化技巧

提升医学图像分割模型性能的实用技巧：

1. 数据层面优化：

   • 使用更丰富的数据增强策略
   • 应用测试时增强(TTA)
   • 平衡类别分布

2. 模型层面优化：

   • 尝试不同的损失函数组合
   • 调整学习率调度策略
   • 使用混合精度训练

3. 后处理优化：

   • 应用形态学操作平滑分割边界
   • 使用条件随机场(CRF)细化结果
   • 设置合适的分割阈值

以下是一个结合了多种优化技巧的训练示例：

from torch.cuda.amp import GradScaler, autocast

def train_with_amp(model, train_loader, val_loader, epochs=100):
    device = torch.device('cuda')
    model = model.to(device)
    
    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)
    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, epochs)
    criterion = DiceBCELoss()
    scaler = GradScaler()
    
    for epoch in range(epochs):
        model.train()
        train_loss = 0
        for images, masks in tqdm(train_loader):
            images, masks = images.to(device), masks.to(device)
            
            optimizer.zero_grad()
            
            # 混合精度训练
            with autocast():
                outputs = model(images)
                loss = criterion(outputs, masks)
            
            scaler.scale(loss).backward()
            scaler.step(optimizer)
            scaler.update()
            
            train_loss += loss.item()
        
        scheduler.step()
        
        # 测试时增强评估
        val_score = evaluate_with_tta(model, val_loader, device)
        
        print(f'Epoch {epoch+1}/{epochs} | Loss: {train_loss/len(train_loader):.4f} | Val Dice: {val_score:.4f}')

7. 常见问题与解决方案

在实际项目中可能会遇到以下典型问题及解决方法：

1. 数据量不足：

   • 使用迁移学习，利用SAM2的预训练权重
   • 应用更强大的数据增强
   • 尝试小样本学习技术

2. 类别不平衡：

   • 使用加权损失函数
   • 采用焦点损失(Focal Loss)
   • 过采样少数类或欠采样多数类

3. 模型过拟合：

   • 增加正则化（Dropout, Weight Decay）
   • 使用早停策略
   • 简化模型结构

4. 推理速度慢：

   • 使用模型量化
   • 尝试知识蒸馏
   • 优化输入分辨率

以下是一个处理类别不平衡的加权损失函数实现：

class WeightedDiceBCELoss(nn.Module):
    def __init__(self, pos_weight=2.0):
        super().__init__()
        self.pos_weight = pos_weight
        
    def forward(self, pred, target):
        pred = torch.sigmoid(pred)
        
        # 计算正负样本权重
        pos = target.sum()
        neg = target.numel() - pos
        weight = torch.where(target > 0, 
                           torch.tensor(neg/pos * self.pos_weight).to(target.device), 
                           torch.tensor(1.0).to(target.device))
        
        # 加权Dice损失
        intersection = (weight * pred * target).sum()
        dice_loss = 1 - (2. * intersection) / ((weight * pred).sum() + (weight * target).sum() + 1e-6)
        
        # 加权BCE损失
        bce = F.binary_cross_entropy_with_logits(pred, target, reduction='none')
        bce = (weight * bce).mean()
        
        return dice_loss + bce

在医疗AI领域，模型的可解释性至关重要。我们可以使用以下方法增强模型透明度：

def visualize_attention(model, image_path):
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    model = model.to(device)
    
    # 加载图像
    image = cv2.imread(image_path)
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    original_size = image.shape[:2]
    
    # 预处理
    transform = A.Compose([
        A.Resize(256, 256),
        A.Normalize()
    ])
    transformed = transform(image=image)
    img_tensor = torch.from_numpy(transformed['image']).permute(2,0,1).unsqueeze(0).float().to(device)
    
    # 注册钩子获取注意力
    attentions = []
    def hook_fn(module, input, output):
        attentions.append(output.cpu())
    
    # 选择目标层（通常是编码器的最后一层）
    target_layer = model.encoder.blocks[-1]
    handle = target_layer.register_forward_hook(hook_fn)
    
    # 前向传播
    with torch.no_grad():
        _ = model(img_tensor)
    
    # 移除钩子
    handle.remove()
    
    # 可视化注意力
    attn = attentions[0].mean(dim=1).squeeze()
    attn = F.interpolate(attn.unsqueeze(0).unsqueeze(0), 
                        size=original_size, 
                        mode='bilinear').squeeze().numpy()
    
    # 叠加到原图
    heatmap = cv2.applyColorMap(np.uint8(255 * attn), cv2.COLORMAP_JET)
    superimposed = cv2.addWeighted(image, 0.5, heatmap, 0.5, 0)
    
    return image, attn, superimposed