3D U-Net在医学图像分割中的原理与实践：从模型架构到性能优化

1次阅读

没有评论

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

医学图像分割是医疗 AI 中的核心任务，在肿瘤检测、器官分割等场景中发挥着重要作用。传统 2D 方法虽然简单，但存在明显局限性：

空间信息丢失 ：将 3D 医学图像（如 CT、MRI）切片处理时，无法建模层间关系
伪影风险 ：逐层预测可能导致器官边界不连续
感受野受限 ：难以捕捉大体积病变的整体特征

3D U-Net 的卷积核增加深度维度（如 3×3×3 vs 2D 的 3×3），参数量呈立方增长。以基础版本为例：

2D U-Net：约 7M 参数
3D U-Net：约 19M 参数

实际测试显示（NVIDIA V100 显卡）：

输入尺寸 128×128×128 时
2D 版本：每秒 30 样本
3D 版本：每秒 8 样本

3D 卷积能同时捕获 XY 平面和 Z 轴特征，对囊状器官（如肾脏）分割更有利

class MedicalDataset(Dataset):
    def __init__(self, scan_dir, mask_dir):
        self.scans = sorted(glob(f"{scan_dir}/*.nii.gz"))  # 读取 NIfTI 格式
        self.masks = sorted(glob(f"{mask_dir}/*.nii.gz"))

    def __getitem__(self, idx):
        scan = nib.load(self.scans[idx]).get_fdata()  # 形状 [D,H,W]
        mask = nib.load(self.masks[idx]).get_fdata()
        return torch.FloatTensor(scan).unsqueeze(0), torch.FloatTensor(mask)  # 增加通道维

class DoubleConv(nn.Module):
    """(3D 卷积 → BN → ReLU) × 2"""
    def __init__(self, in_ch, out_ch):
        super().__init__()
        self.conv = nn.Sequential(nn.Conv3d(in_ch, out_ch, kernel_size=3, padding=1),
            nn.BatchNorm3d(out_ch),
            nn.ReLU(inplace=True),
            nn.Conv3d(out_ch, out_ch, kernel_size=3, padding=1),
            nn.BatchNorm3d(out_ch),
            nn.ReLU(inplace=True)
        )

class Down(nn.Module):
    """下采样层: MaxPool + DoubleConv"""
    def __init__(self, in_ch, out_ch):
        super().__init__()
        self.mpconv = nn.Sequential(nn.MaxPool3d(2),
            DoubleConv(in_ch, out_ch)
        )

弹性变形 ：模拟器官自然形变

from torchio.transforms import RandomElasticDeformation
transform = RandomElasticDeformation(
    num_control_points=7,
    max_displacement=15
)

模拟伪影 ：增加运动伪影和金属伪影

def hybrid_loss(pred, target):
    bce = F.binary_cross_entropy_with_logits(pred, target)
    pred = torch.sigmoid(pred)
    dice = 1 - (2.*(pred*target).sum() + 1e-7) / (pred.sum() + target.sum() + 1e-7)
    return 0.5*bce + 0.5*dice

CT 值截断范围建议：