3D医学图像分割网络入门实战：从数据预处理到模型训练全流程解析

1次阅读

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

医学图像分割是 AI 辅助诊断的重要基础任务，在肿瘤定位、手术规划等场景有不可替代的价值。相比自然图像，CT/MRI 等 3D 医学数据具有两个显著特点：

各向异性分辨率：层内分辨率（如 512×512）通常远高于层间分辨率（如 2mm 层厚），直接导致三维卷积核感受野失衡
标注成本极高：专家标注单个病例常需数小时，且不同机构标注标准不一致，引发标签噪声问题

2D 分割网络（如 UNet）
优势：显存占用低，可直接借用自然图像领域的预训练权重
劣势：无法捕捉层间上下文，对连续器官（如血管）分割效果差
3D 分割网络（如 UNet3D/VNet）
优势：保持空间一致性，对复杂结构建模能力更强
劣势：计算复杂度呈立方增长，需特殊优化策略

经典架构对比：

网络结构	参数量	适用场景
UNet3D	约 19M	中等规模器官（肝脏等）
VNet	约 65M	大体积靶区（前列腺等）

import nibabel as nib

def load_nii(path):
    img = nib.load(path)
    data = img.get_fdata()
    affine = img.affine  # 保存空间坐标信息
    return np.ascontiguousarray(data)

截断无效值：data = np.clip(data, -1000, 1000)
器官特定标准化：
肝脏窗宽：(data - 40) / 160
肺窗宽：(data + 1000) / 1400

class PatchSampler:
    def __init__(self, vol_shape, patch_size=(128,128,32)):
        self.strides = [s//4 for s in patch_size]  # 75% 重叠

    def __call__(self, volume):
        patches = []
        for z in range(0, depth, self.strides[2]):
            # 类似处理 x,y 维度...
            patch = volume[z:z+patch_size[2]]
            patches.append(patch)
        return patches

关键操作：

随机弹性变形
轴向镜像翻转
±10% 尺度抖动

复合损失函数公式：

$$
\mathcal{L} = 0.5\cdot\text{Dice} + 0.5\cdot\text{CE} + \lambda\cdot\text{边界损失}
$$

类别权重计算：

class_weight = 1 / (np.bincount(label.flatten()) + 1e-6)

class MedicalDataset(Dataset):
    def __init__(self, img_paths, label_paths):
        self.img_paths = img_paths
        self.transform = Compose([RandomRotate90(p=0.5),
            GaussianNoise(p=0.2)
        ])

    def __getitem__(self, idx):
        img = load_nii(self.img_paths[idx])
        img = self.hu_window(img, organ='liver')

        if self.label_paths:
            label = load_nii(self.label_paths[idx])
            img, label = self.transform(img, label)
            return torch.FloatTensor(img), torch.LongTensor(label)
        return torch.FloatTensor(img)

class UNet3D(nn.Module):
    def __init__(self, in_channels=1):
        super().__init__()
        self.encoder1 = nn.Sequential(nn.Conv3d(in_channels, 32, 3, padding=1),
            nn.BatchNorm3d(32),
            nn.ReLU())
        # 下采样层...

    def forward(self, x):
        if x.dim() == 4:  # 单模态
            x = x.unsqueeze(1)
        x1 = self.encoder1(x)
        # 编解码结构...
        return x

使用 DistributedDataParallel 而非DataParallel
BatchNorm 层替换为SyncBatchNorm
验证阶段关闭梯度同步

def predict_large_volume(model, volume, patch_size):
    output = torch.zeros_like(volume)
    counts = torch.zeros_like(volume)

    for patch, coord in PatchSampler(patch_size)(volume):
        pred = model(patch)
        output[coord] += pred
        counts[coord] += 1

    return output / counts.clamp(min=1)  # 重叠区取平均