3DUNet医学图像分割实战：肝脏肿瘤分割源码解析与优化指南

1次阅读

没有评论

共计 2528 个字符，预计需要花费 7 分钟才能阅读完成。

医学图像分割是 AI 辅助诊断中的核心任务，尤其在肝脏肿瘤诊断中，精准分割能帮助医生量化肿瘤体积、评估治疗效果。但实际操作中开发者常遇到两大挑战：

数据特性问题：肝脏与肿瘤的边界模糊（尤其是 HCC 肝癌），CT 图像中不同组织间灰度值重叠严重
资源限制：3D 医学影像（如 512×512×300 的 CT 体积）显存占用大，直接训练全分辨率模型需专业级 GPU

2D CNN：切片级处理，计算效率高但丢失空间上下文信息，对肝脏这类连续器官分割效果有限
3D CNN：体素级处理，能捕捉各向异性特征（如层厚 2.5mm vs 像素 0.7mm），但计算复杂度呈立方增长

相比 V -Net 等方案，3DUNet 在医疗影像中表现更优：

跳跃连接 (Skip Connection)：编码器(encoder) 的多尺度特征与解码器 (decoder) 融合，有效恢复空间细节
深度监督：中间层输出辅助损失函数，缓解梯度消失
内存效率：通过 patch-based 训练策略，在消费级 GPU（如 RTX 3090）上可处理 128×128×64 的输入块

# 医学影像特有的窗宽窗位调整（Liver 窗：WW=150, WL=30）def apply_window(image, window_center=30, window_width=150):
    min_val = window_center - window_width // 2
    max_val = window_center + window_width // 2
    image = np.clip(image, min_val, max_val)
    return (image - min_val) / (max_val - min_val)

# 各向异性重采样至 1mm³体素
from scipy.ndimage import zoom
resampled_data = zoom(original_data, 
                     (original_spacing[0]/1.0, 
                      original_spacing[1]/1.0,
                      original_spacing[2]/1.0))

class DoubleConv(nn.Module):
    """(Conv3D -> BN -> ReLU) × 2"""
    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.double_conv = nn.Sequential(nn.Conv3d(in_channels, out_channels, kernel_size=3, padding=1),
            nn.BatchNorm3d(out_channels),
            nn.ReLU(inplace=True),
            nn.Conv3d(out_channels, out_channels, kernel_size=3, padding=1),
            nn.BatchNorm3d(out_channels),
            nn.ReLU(inplace=True)
        )

    def forward(self, x):
        return self.double_conv(x)

class DownSample(nn.Module):
    """MaxPool3D + DoubleConv"""
    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.maxpool_conv = nn.Sequential(nn.MaxPool3d(2),
            DoubleConv(in_channels, out_channels)
        )

    def forward(self, x):
        return self.maxpool_conv(x)

from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()

with autocast():
    output = model(input)
    loss = criterion(output, target)

scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()

def sliding_window_inference(inputs, model, roi_size=128, sw_batch_size=4):
    """
    roi_size: 子区域大小
    sw_batch_size: 并行处理的子区域数
    """
    return sliding_window_inference(inputs, roi_size, sw_batch_size, model, overlap=0.5, mode="gaussian")

损失函数调整：组合 Dice Loss（促进区域重叠）与 Focal Loss（关注难样本）

class DiceFocalLoss(nn.Module):
    def __init__(self, gamma=2.0):
        super().__init__()
        self.gamma = gamma

    def forward(self, pred, target):
        # Dice term
        smooth = 1e-5
        intersection = (pred * target).sum()
        dice = (2. * intersection + smooth) / (pred.sum() + target.sum() + smooth)

        # Focal term
        bce = F.binary_cross_entropy(pred, target, reduction='none')
        focal = (1. - torch.exp(-bce)) ** self.gamma * bce

        return 1 - dice + focal.mean()