3D U-Net医学图像分割实战：从零实现肝脏肿瘤分割代码

1次阅读

没有评论

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

医学图像分割在临床诊断和治疗规划中扮演着关键角色。以肝脏肿瘤为例，精确的分割结果可以帮助医生：

准确测量肿瘤体积，评估疾病进展
制定个性化的手术或放疗方案
减少人为判读的主观性和误差

然而，医学图像分割也面临独特挑战：

数据维度高：3D 医学影像（如 CT/MRI）包含数百层切片
标注成本高：需要专业放射科医生参与
类别不平衡：肿瘤区域可能只占图像的极小部分
隐私合规：医疗数据需符合 HIPAA 等法规要求

逐层处理 2D 切片
计算资源需求低
可能丢失层间上下文信息

直接处理 3D 体数据
保留空间连续性
更适合肝脏等复杂器官

# 典型 3D U-Net 输入输出尺寸示例
input_shape = (128, 128, 128, 1)  # (D,H,W,C)
output_shape = (128, 128, 128, 2) # 背景 + 肿瘤两类

import nibabel as nib
import numpy as np

# 加载 NIfTI 文件
def load_nii(filepath):
    scan = nib.load(filepath)
    data = scan.get_fdata()
    # 归一化到 [0,1]
    data = (data - np.min(data)) / (np.max(data) - np.min(data))
    return np.expand_dims(data, axis=-1)  # 添加通道维度

import torch
import torch.nn as nn

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)
        )

class UNet3D(nn.Module):
    def __init__(self, n_classes=2):
        super(UNet3D, self).__init__()
        # 编码器部分
        self.encoder1 = DoubleConv(1, 64)
        self.encoder2 = DoubleConv(64, 128)
        # 解码器部分
        self.upconv1 = nn.ConvTranspose3d(128, 64, kernel_size=2, stride=2)
        self.decoder1 = DoubleConv(128, 64)
        # 最终输出层
        self.final_conv = nn.Conv3d(64, n_classes, kernel_size=1)

class DiceBCELoss(nn.Module):
    def __init__(self, smooth=1.):
        super(DiceBCELoss, self).__init__()
        self.smooth = smooth

    def forward(self, inputs, targets):
        # Dice 系数计算
        inputs = torch.sigmoid(inputs)
        intersection = (inputs * targets).sum()
        dice_loss = 1 - (2.*intersection + self.smooth)/
                    (inputs.sum() + targets.sum() + self.smooth)
        # BCE 损失
        bce = nn.functional.binary_cross_entropy(inputs, targets)
        return dice_loss + bce

from torchvision.transforms import Compose

# 3D 专用增强变换
transforms = Compose([RandomRotate3D(angles=[0, 90, 180]),
    RandomFlip3D(),
    ElasticDeformation3D(),
    AddGaussianNoise(mean=0, std=0.01)
])

model = UNet3D()
if torch.cuda.device_count() > 1:
    print(f"Using {torch.cuda.device_count()} GPUs!")
    model = nn.DataParallel(model)
model.to(device)

def dice_coef(y_true, y_pred):
    intersection = np.sum(y_true * y_pred)
    return (2. * intersection) / (np.sum(y_true) + np.sum(y_pred))

def hausdorff_distance(y_true, y_pred):
    # 使用 scipy 计算 95% Hausdorff 距离
    from scipy.spatial.distance import directed_hausdorff
    return max(directed_hausdorff(y_true, y_pred)[0],
               directed_hausdorff(y_pred, y_true)[0])

使用梯度累积：

optimizer.zero_grad()
for _ in range(accum_steps):
    outputs = model(inputs)
    loss = criterion(outputs, labels) / accum_steps
    loss.backward()
optimizer.step()