BRATS医学图像分割实战：从数据预处理到3D U-Net模型优化

1次阅读

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

BRATS（Brain Tumor Segmentation）数据集是脑肿瘤分割领域的标杆数据集，包含多模态 MRI（T1/T1c/T2/FLAIR）的 3D 扫描数据，标注了肿瘤核心（enhancing tumor）、水肿区域（edema）和坏死组织（necrosis）等子结构。这个数据集对推动 AI 辅助诊断有三大核心价值：

多模态数据提供了互补的解剖和功能信息（如 T1c 增强序列对肿瘤边界的敏感性）
标注由多名专家共同完成，具有临床参考意义
每年更新的比赛机制持续推动算法进步

但在实际应用中，开发者会面临三个典型挑战：

多模态配准（multi-modal registration）问题：不同序列的扫描存在空间偏移
类别不平衡（class imbalance）：正常脑组织占比远大于肿瘤区域
小样本训练：标注数据量有限（BraTS 2021 仅含 1251 例训练数据）

使用 SimpleITK 加载 NIfTI 格式数据时，需特别注意各向异性（anisotropy）体素（voxel）间距的处理：

import SimpleITK as sitk
import numpy as np

def load_nii_to_tensor(file_path):
    # 读取时保留原始间距信息
    sitk_image = sitk.ReadImage(file_path)
    spacing = sitk_image.GetSpacing()  # 获取体素物理间距(mm)

    # 转换为 numpy 数组并调整轴顺序
    data = sitk.GetArrayFromImage(sitk_image)  # 形状为(D,H,W)
    data = np.transpose(data, (2, 1, 0))  # 调整为(W,H,D)

    # 窗宽窗位调整（以 FLAIR 序列为例）if 'flair' in file_path.lower():
        data = np.clip(data, -0, 500)  # 去除极端值
        data = (data - data.min()) / (data.max() - data.min())

    return torch.FloatTensor(data), spacing

核心改进点在于 3D 卷积（3D convolution）和跳跃连接（skip connection）的实现：

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 UNet3D(nn.Module):
    def __init__(self, n_classes=4):
        super().__init__()
        # 下采样路径
        self.down1 = DoubleConv(4, 32)  # 4 模态输入
        self.down2 = DoubleConv(32, 64)
        # 上采样路径
        self.up1 = nn.ConvTranspose3d(64, 32, kernel_size=2, stride=2)
        self.conv1 = DoubleConv(64, 32)  # skip connection 拼接后通道数
        # 最终输出层
        self.out = nn.Conv3d(32, n_classes, kernel_size=1)

    def forward(self, x):
        # 下采样阶段
        x1 = self.down1(x)
        x2 = F.max_pool3d(x1, 2)
        x2 = self.down2(x2)

        # 上采样阶段
        x = self.up1(x2)
        x = torch.cat([x, x1], dim=1)  # skip connection
        x = self.conv1(x)

        return self.out(x)

针对小区域分割，采用 Dice Loss（戴斯损失）与 CrossEntropy（交叉熵）的加权组合：

def dice_loss(pred, target, smooth=1e-5):
    # pred: [B, C, D, H, W]
    # target: [B, D, H, W]
    pred = F.softmax(pred, dim=1)
    target_onehot = F.one_hot(target, num_classes=4).permute(0,4,1,2,3)

    intersection = (pred * target_onehot).sum(dim=(2,3,4))
    union = pred.sum(dim=(2,3,4)) + target_onehot.sum(dim=(2,3,4))

    # 对 EDEMA/NECROTIC 类别赋予更高权重
    class_weights = torch.tensor([0.1, 0.3, 0.3, 0.3]).to(pred.device)
    dice = (2. * intersection + smooth) / (union + smooth)
    return 1 - (class_weights * dice).sum()

使用 PyTorch 的 DistributedDataParallel 时需注意：

每个进程需独立加载数据片段
Batch size 指单卡处理的样本量
验证集需在主进程集中计算

# 初始化进程组
torch.distributed.init_process_group(backend='nccl')

# 模型包装
model = nn.SyncBatchNorm.convert_sync_batchnorm(model)
model = DDP(model, device_ids=[local_rank])

# 数据采样器
sampler = DistributedSampler(dataset, shuffle=True)
dataloader = DataLoader(dataset, batch_size=8, sampler=sampler)

测试大体积数据时，采用滑动窗口（sliding window）策略：