3DUNet医学图像分割实战：从数据预处理到模型部署的全流程指南

1次阅读

没有评论

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

医学图像分割是医疗 AI 中一个非常重要的任务，但新手入门时往往会遇到几个主要问题：

数据方面：医学影像通常样本量小，标注成本高，而且标注质量参差不齐
计算资源：3D 医学影像数据量大，显存占用高，训练过程对硬件要求高
部署落地：在临床环境中，模型需要在边缘设备上高效运行，对推理速度有严格要求

在医学图像分割领域，常见的有几种架构选择：

2D CNN：
优点：计算量小，实现简单
缺点：无法捕捉切片间的空间信息
3D CNN：
优点：能完整利用三维空间信息
缺点：计算复杂度高
Transformer：
优点：长距离依赖建模能力强
缺点：需要大量数据，计算开销大

综合考虑医学影像的特点（数据量小但空间信息重要），3DUNet 是一个很好的平衡点。它继承了 UNet 的优秀特性，同时通过 3D 卷积更好地处理体积数据。

医学影像通常以 DICOM 格式存储，我们需要先转换为更适合处理的 NIfTI 格式。使用 dicom2nifti 库可以轻松完成这个转换：

import dicom2nifti

dicom2nifti.convert_directory('input_dicom', 'output_nifti')

CT 图像的窗宽 (Window Width) 和窗位 (Window Center) 调整非常重要，这相当于图像的对比度和亮度调节。以下是示例代码：

import numpy as np

def apply_window(image, window_center, window_width):
    min_val = window_center - window_width / 2
    max_val = window_center + window_width / 2
    windowed = np.clip(image, min_val, max_val)
    return (windowed - min_val) / (max_val - min_val)

下面是一个改进版 3DUNet 的核心架构代码，主要特点是加入了更深层的跳跃连接：

import torch
import torch.nn as nn

class DoubleConv(nn.Module):
    """(convolution => [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)

# 完整 UNet 架构代码较长，此处省略...

使用混合精度训练可以显著减少显存占用并加快训练速度：

from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()

with autocast():
    outputs = model(inputs)
    loss = criterion(outputs, labels)

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

当显存不足时，可以通过梯度累积来模拟更大的 batch size：

accumulation_steps = 4

for i, (inputs, labels) in enumerate(train_loader):
    with autocast():
        outputs = model(inputs)
        loss = criterion(outputs, labels) / accumulation_steps

    scaler.scale(loss).backward()

    if (i + 1) % accumulation_steps == 0:
        scaler.step(optimizer)
        scaler.update()
        optimizer.zero_grad()

不同模态的医学影像需要不同的归一化方式：

模态	归一化方法
CT	固定 HU 值范围 (-1000 到 3000)
MRI	各向同性归一化 (减去均值除以标准差)

标准的 Dice Loss 在处理边缘模糊的器官时效果不佳，可以尝试：

class ImprovedDiceLoss(nn.Module):
    def __init__(self, smooth=1e-5):
        super().__init__()
        self.smooth = smooth

    def forward(self, pred, target):
        # 加入边缘权重
        edge_weight = compute_edge_weight(target)
        intersection = (pred * target * edge_weight).sum()
        union = (pred + target).sum()
        return 1 - (2. * intersection + self.smooth) / (union + self.smooth)

在 BraTS 2020 数据集上的测试结果：