简介:面向医学图像分割的实战资源,基于DenseUnet完成腹部多脏器(背景、肝脏、右肾、左肾、脾脏)的像素级分割,适合有一定深度学习基础、希望复现完整训练流程并掌握分割项目组织方式的开发者。资源共31个文件,压缩包约347MB,包括8个Python脚本(训练、验证、预测、混淆矩阵等)、2个模型权重文件、训练日志与可视化曲线(loss/IoU、学习率衰减、数据集可视化)、图像示例、XML配置及README说明。代码注释详细,训练、验证、预测三个模块分工明确:训练脚本自动生成训练/验证集的loss、IoU曲线及学习率衰减曲线;验证脚本计算测试集的IoU、recall、precision、像素准确率;预测脚本生成ground truth及叠加掩膜图像。模型训练50个epoch,测试集像素准确率达0.987,平均IoU达0.798。已有166人学习,适合作为多器官分割任务的参考实现与基线框架。
1. 基于 DenseUnet 的腹部多脏器分割:不是把 U-Net 加粗一层就能交差
医学图像分割任务里,腹部多脏器分割和单器官分割的难度完全不在一个数量级。单器官可以只盯着肝脏或者肾脏的轮廓,多脏器却要在一张腹部 CT 里同时把肝脏、脾脏、左右肾、胰腺甚至胆囊分开,器官之间的灰度重叠、边界粘连、尺寸悬殊,会让很多常见分割模型直接翻车。这也是很多团队从 U-Net、Res-UNet 一路换到 DenseUnet 的原因:DenseUnet 用密集连接把浅层纹理和深层语义缝在一起,在训练数据量不够大、标注边界又不干净的真实项目里,收敛更稳,小器官也不会一上来就被大器官吞掉。这篇笔记面向正在做医学影像分割、手上已经有一批腹部 CT 或公开数据集的人,从数据预处理、模型结构、损失函数到训练排错,按完整落地顺序讲一遍,代码可以直接改去跑自己的数据。
2. 为什么选 DenseUnet 做腹部多脏器:密集连接带来的取舍
2.1 DenseBlock 到底改了什么
DenseNet 最核心的改变不是网络变深,而是每一层都把前面所有层的特征图拼接起来作为输入。假设一个 DenseBlock 里第 i 层输入通道数是 k0 + (i-1) * growth_rate,输出的特征图会拼接到后面的每一层,而不是像 ResNet 那样只做一次加法。growth_rate 一般取 12、16 或者 24,它控制每层新增多少通道,通道数不会像普通卷积网络那样逐层爆炸。对腹部 CT 这种标注样本普遍只有几十到一两百例的场景,参数总量小、特征复用度高,训练集不容易过拟合,这是选择 DenseUnet 最直接的理由。
从梯度角度看,DenseBlock 相当于给反向传播铺了多条近路。锁骨下的深层卷积层哪怕处在网络最后,梯度也能通过密集连接直达前层,缓解深层网络常见的梯度消失问题。很多人在 U-Net 上加深 encoder,从 4 层加到 6 层后训练 loss 却不降,换成 DenseBlock 之后反而能稳定下降,就是这个原因。代价是显存:特征图被反复拼接,前向传播时要保存大量中间张量,DenseUnet 的显存开销比同参数量 ResNet 编码器高不少。实际项目里我一般把 growth_rate 压到 12 或 16,而不是照搬论文里的 32。
2.2 U 形结构与 DenseUnet 的融合方式
DenseUnet 并不是简单地在 U-Net 里把卷积块换成 DenseBlock。常见做法是用 4 个 DenseBlock 构成编码器,每个 Block 后面接一个 Transition Down 做下采样,下采样包括 1x1 卷积降通道和 2x2 平均池化。解码器保持 U-Net 的对称结构,转置卷积把空间尺寸恢复回来,再通过 skip connection 把编码器对应层的特征拼接上。这里的拼接不是短接,而是 channel 维度上的 concat,解码器能得到编码器每层保留的细节。
因为腹部多脏器的标签本质上是逐像素的多分类,也就是医学语义分割里的标准任务,和实例分割这种要把每个单独器官实例区分开的问题不一样。DenseUnet 的输出层直接接一个 1x1 卷积,把特征图映射到器官类别数,配合 softmax 得到每个像素属于每个器官的概率。如果要做的是把同一类别的多个肿瘤分开,那就要转向实例分割的思路,DenseUnet 的标签和损失函数都得换,不要混为一谈。
2.3 为什么多脏器任务比单器官更吃这种结构
腹部多脏器分割的难点集中在三点:器官灰度范围重叠、大小悬殊、边界互相挤压。肝脏和脾脏在 CT 上都是软组织密度,单纯靠窗宽窗位很难从灰度上把两者分开;胰腺又小又薄,常被胃、十二指肠和血管包夹;肾上腺体积极小,普通模型经常把它整个漏掉。DenseBlock 的多层特征拼接,相当于把不同感受野的信息同时摆在解码器面前,大器官靠深层语义锁定位置,小器官靠浅层边缘纹理保住存在感,比单一路径的卷积更有优势。
另外,医学图像分割训练集通常很小,常见做法是拿公开的腹部多器官数据集或者医院积累几十例真实病例做训练。这种规模下,模型很容易在某一类器官上过分自信,验证集上大器官分数很高,小器官却趋近于零。DenseUnet 的特征重用相当于给所有类别共享一组更紧凑的底层特征,降低了小器官类别的学习门槛。当然它不是万能的,如果数据集本身存在严重标注错位,再密集的连接也补不回来,这一点在后面的排查章节里会展开讲。
3. 数据集准备与预处理:从 nii 到能喂给 DenseUnet 的张量
3.1 拿到数据先做三件事:读、看、查
不管数据是从公开渠道下载的还是医院影像科导出的,第一步都别急着训练。先确认三个东西:图像格式、spacing 方向、标签类别数。腹部分割数据最常见的格式是 NIfTI,也就是 .nii.gz 文件,一个文件包含三维体数据和一个仿射矩阵。用下面的代码读取并打印基本信息:
import nibabel as nib import numpy as np def inspect_nii(path): img = nib.load(path) data = np.asanyarray(img.dataobj) affine = img.affine spacing = np.sqrt((affine[:3, :3] ** 2).sum(axis=0)) print(f"shape: {data.shape}") print(f"spacing: {spacing}") print(f"data min/max: {data.min():.1f} / {data.max():.1f}") labels = np.unique(data) print(f"label values: {labels}") return data, spacing打印 shape 是为了确认三维体数据的排列顺序,不同工具导出的 NIfTI 可能把轴序放在切片维度上不一致。打印 spacing 是为了后面重采样,因为不同医院扫描层厚可能不一样,有的是 1.5 mm,有的是 5 mm,如果直接混着训练,模型会去学扫描参数而不是解剖结构。打印 label values 是最重要的一步,很多数据集转出来标签并不是从 0 开始连续排列的,可能背景是 0,肝脏是 1,脾脏直接跳成 11,如果不做类别映射,训练时类别数对不上,loss 直接报错。
3.2 重采样、窗宽窗位、归一化
腹部 CT 的原始数值是 HU,范围通常在 -1024 到 3000 以上,直接送进网络训练会出问题。常规做法是先做重采样,把不同层厚的数据统一到一个固定 spacing,我一般用 1.5x1.5x3.0 mm,这样既能保留大部分解剖细节,又不会让数据量过大。这里用 SimpleITK 处理比较省事:
import SimpleITK as sitk def resample_to_spacing(image_path, target_spacing=(1.5, 1.5, 3.0)): img = sitk.ReadImage(image_path) original_spacing = img.GetSpacing() original_size = img.GetSize() new_spacing = target_spacing new_size = [ int(round(original_size[0] * original_spacing[0] / new_spacing[0])), int(round(original_size[1] * original_spacing[1] / new_spacing[1])), int(round(original_size[2] * original_spacing[2] / new_spacing[2])), ] resampler = sitk.ResampleImageFilter() resampler.SetOutputSpacing(new_spacing) resampler.SetSize(new_size) resampler.SetOutputDirection(img.GetDirection()) resampler.SetOutputOrigin(img.GetOrigin()) resampler.SetInterpolator(sitk.sitkLinear) return resampler.Execute(img), new_size注意,图像用线性插值,标签掩膜绝不能用线性插值,要用最近邻插值,否则器官边缘会出现本来不存在的灰度值,比如标签 1 和标签 3 之间插出个 2,类别就乱了。重采样之后做窗宽窗位处理。腹部多脏器没有单一最优窗,我一般取 [-125, 250] HU,把软组织、肝脏、脾脏、肾脏都包进来,胰腺也能保留一定对比度。然后归一化到 [-1, 1] 或 [0, 1] 区间,方便模型收敛:
HU_MIN, HU_MAX = -125.0, 250.0 def preprocess_ct(volume): volume = np.clip(volume, HU_MIN, HU_MAX) volume = (volume - HU_MIN) / (HU_MAX - HU_MIN) # [0, 1] volume = volume * 2.0 - 1.0 # [-1, 1] return volume.astype(np.float32)窗宽窗位的选择对训练结果影响很大。窗太窄会把肝内血管和正常组织压成一团,窗太宽又会让低对比度的胰腺变得几乎不可见。如果只做一个器官,比如肝脏,可以取 [-100, 200],让肝实质更亮;但多脏器必须兼顾,取值别太极端。
3.3 数据划分:按患者切,绝不按切片切
腹部 CT 一个病例就是几十张甚至几百张切片,如果直接把所有切片打乱后随机分训练集和验证集,同一个患者的切片会同时出现在两边,这叫数据泄漏。验证集指标会虚高,等到部署在新患者身上立刻现出原形。正确做法是把病例 ID 作为分组单位,先按患者分组,再整体划分:
import os import numpy as np import json all_cases = sorted(os.listdir("your_ct_folder")) np.random.seed(42) np.random.shuffle(all_cases) split_idx = int(len(all_cases) * 0.8) train_cases = all_cases[:split_idx] val_cases = all_cases[split_idx:] with open("train_cases.json", "w") as f: json.dump(train_cases, f, indent=2)我一般会把训练集、验证集的病例名单存成 JSON,而不是只记录切片的数量。因为后面每次实验都要保证同一个病例落在同一边,固定名单才能让指标有可比性。另外,建议留 5% 到 10% 的病例做冻结测试集,完全不参与调参,只在最终模型上跑一次,这是医学影像项目里容易被忽略的细节。
4. 训练与评估:用 PyTorch 把 DenseUnet 跑通
4.1 核心网络结构:DenseBlock 与 U 形解码器
在 2D 模式下,输入是一张腹部 CT 轴向切片,形状是 [B, 1, H, W],标签是 [B, H, W],每个像素的值是器官类别 ID。PyTorch 里 DenseBlock 的核心代码如下,实际工程里可以把 3D 卷积换成 2D,其余结构几乎不动:
import torch import torch.nn as nn class DenseLayer(nn.Module): def __init__(self, in_ch, growth_rate, bn_size=4, drop_rate=0.2): super().__init__() mid_ch = bn_size * growth_rate self.block = nn.Sequential( nn.BatchNorm2d(in_ch), nn.ReLU(inplace=True), nn.Conv2d(in_ch, mid_ch, kernel_size=1, bias=False), nn.BatchNorm2d(mid_ch), nn.ReLU(inplace=True), nn.Conv2d(mid_ch, growth_rate, kernel_size=3, padding=1, bias=False), ) self.drop_rate = drop_rate def forward(self, x): out = self.block(x) if self.drop_rate > 0: out = nn.functional.dropout2d(out, p=self.drop_rate, training=self.training) return torch.cat([x, out], dim=1) class DenseBlock(nn.Module): def __init__(self, in_ch, num_layers, growth_rate=12): super().__init__() layers = [] ch = in_ch for _ in range(num_layers): layers.append(DenseLayer(ch, growth_rate)) ch += growth_rate self.layers = nn.Sequential(*layers) def forward(self, x): return self.layers(x)DenseLayer 里先做 1x1 卷积降通道,再做 3x3 卷积输出 growth_rate 个特征图。bn_size=4 控制 1x1 卷积的中间通道数,growth_rate=12 每次只新增 12 个通道,整个 DenseBlock 的通道增长是线性的。Dropout 在 DenseBlock 里很关键,因为密集连接会特征重用,不加 dropout 在小数据集上非常容易过拟合。inplace=True 能省一点显存,但开启后如果模型里的张量被其他地方引用,反向传播可能出问题,显存不够时再开。
解码器部分采用 U-Net 结构,编码器输出再经过转置卷积上采样,每一层把编码器对应层的特征在通道维度上拼接,然后接两个 3x3 卷积。最后用 1x1 卷积把通道数压成器官类别数。类别数包括背景,如果任务是分割肝脏、脾脏、肾脏、胰腺四类器官,n_classes 就是 5。DenseBlock 下采样之间用 Transition 层把通道数压缩到原来的 0.5 倍,否则拼接多了显存扛不住。
4.2 损失函数:Dice Loss 和 CrossEntropy 的搭配
腹部多脏器分割的标签类别极不平衡。肝脏在轴位切片上可能占 40% 的面积,肾上腺可能只占 0.5%,如果直接用 CrossEntropy,模型会倾向于把所有像素都预测成背景和大器官。我一般用 Soft Dice Loss 和带类别权重的 CrossEntropy 加权组合:
def soft_dice_loss(pred, target, class_weights=None, epsilon=1e-6): # pred: [B, C, H, W] after softmax # target: [B, H, W] B, C, H, W = pred.shape target_one_hot = nn.functional.one_hot(target, num_classes=C).permute(0, 3, 1, 2).float() intersection = (pred * target_one_hot).sum(dim=(0, 2, 3)) union = pred.sum(dim=(0, 2, 3)) + target_one_hot.sum(dim=(0, 2, 3)) dice = (2.0 * intersection + epsilon) / (union + epsilon) if class_weights is not None: dice = dice * class_weights return 1.0 - dice.mean()Dice Loss 对类别不平衡不敏感,因为它直接优化每个类别的重叠度,而不是像素总数。但纯 Dice Loss 在训练初期梯度不稳定,尤其是某个小器官完全没有预测出来时,分母很小,梯度震荡很厉害。常见做法是把 Dice Loss 和加权 CrossEntropy 按 1:1 或 0.8:0.2 混合。类别权重可以按体素占比的倒数归一化得到,比如 background 权重设 0.1,胰腺权重设 5.0,但不要设得过于极端,权重太大会让背景区域的噪声也被放大。
4.3 训练循环、学习率与模型保存
训练时输入切片大小我建议 256x256,batch size 在 8 到 16 之间,优化器用 AdamW 或 SGD。AdamW 收敛快,SGD 需要更多 epoch 但最终表现可能更稳定。学习率初始值 1e-4,配合 ReduceLROnPlateau 在验证指标停滞时降低学习率,比固定学习率省心。验证指标用平均 Dice,也就是所有器官类别的 Dice 取平均,而不是只看肝脏这种大器官。
模型保存的关键是每次验证只保留最优权重。很多人训练完只保存最后一轮,结果最后十轮已经过拟合,白白浪费计算资源。训练伪代码如下:
best_dice = 0.0 for epoch in range(num_epochs): model.train() for ct, label in train_loader: ct = ct.to(device) label = label.to(device) pred = model(ct) loss = mixed_loss(pred, label) optimizer.zero_grad() loss.backward() optimizer.step() val_dice = evaluate(model, val_loader, device) if val_dice > best_dice: best_dice = val_dice torch.save(model.state_dict(), "best_denseunet.pth") scheduler.step(val_dice)我一般把 epoch 数设成 100 到 150,每 10 个 epoch 打印一次各个器官的 Dice。如果到第 20 轮时肝脏已经到 0.9,但胰腺还是 0.1,不要急着调模型,先看数据增强和标签有没有覆盖到胰腺区域,这个在下一章展开。
5. 训练结果不理想的避坑排查:五个真实踩坑记录
5.1 训练 loss 直接 NaN
现象:训练前几个 step 的 loss 是正常浮点数,某个 batch 突然变成 NaN,之后一直 NaN。
原因:最常见的是标签里出现了类别数之外的像素值,比如类别数设置 5,结果标签里有 255;也可能是数据归一化出现无穷大,CT 原始值里有极端值,预处理没做 clip;或者是空切片,整个标签全是背景,CrossEntropy 的类别权重里某个非背景类的分母是零。
解决:训练前扫描所有标签的像素分布,确认np.unique(label)的集合严格落在[0, n_classes-1]范围内;预处理加 clip;在 loss 计算前对 logits 做一个 clamp,比如torch.clamp(logits, min=-100, max=100)。空切片可以直接跳过,也可以在 DataLoader 里过滤掉标签信息量过低的切片。
5.2 训练集 Dice 很高,验证集突然塌掉
现象:训练第 80 轮时 Dice 0.95,验证集只有 0.3,模型像是在验证集上失灵。
原因:最典型的就是第 3.3 节说的数据划分按切片而不是按患者。另一个常见原因是数据增强过度,比如训练时随机裁剪、翻转、弹性形变,但验证时只做了简单中心裁剪,模型对增强后的噪声产生了依赖。
解决:先检查训练集和验证集的 patient ID 是否有交集,交集为零再往下查。数据增强不要全部开满,我一般只开随机水平翻转、轻微旋转和随机裁剪,弹性形变在标注边缘不干净的数据集上慎用。验证时和训练时的预处理必须保持一致,包括窗宽窗位和归一化区间。
5.3 胰腺、肾上腺这种小器官完全分不出来
现象:大器官 Dice 0.9,胰腺或胆囊 Dice 0,整个器官被预测成背景或者被旁边的胃肠区域吞掉。
原因:类别极度不平衡之外,还有一个容易被忽略的问题是下采样次数太多。DenseUnet 编码器如果下采样到原图 1/32,小器官在深层特征图里只有几个像素甚至不到一个像素,解码器再怎么上采样也补不回空间细节。
解决:把下采样倍数从 32 降到 16,最后一层 DenseBlock 输出的特征图分辨率高一点;或者对小器官区域做裁剪训练,单独用胰腺附近的 patch 微调几轮。另外,在损失计算时对小器官类别提高权重,但权重要平滑调整,别把 loss 变成小器官的独角戏。
5.4 分割结果在器官边界出现一圈错误连接
现象:预测结果里,肝和脾之间有一条带状伪影,或者肾脏边缘多出一圈和周围肌肉粘连的区域。
原因:下采样和转置卷积的棋盘效应是一部分,更常见的是训练数据里器官边界标注不干净,不同标注员画的边界差几层像素,模型把这种不干净边界当作特征学到了。
解决:标签后处理时对器官内部做连通域分析,删除面积过小的孤立区域;训练时对标签的边缘做轻微腐蚀,让模型不要过度依赖夸张的边缘细节。验证时也要注意,Dice 对边缘偏移比较敏感,可以叠加 Hausdorff 距离评估边缘质量。
5.5 显存不够,batch size 只能设到 2
现象:一用 DenseUnet 显存就爆,batch size 降到 2 才能跑,训练非常慢。
原因:DenseBlock 前向要保存大量拼接张量,通道重用的代价就是显存。这也是把 growth_rate 设成 48 或 64 时会立刻翻车的直接原因。
解决:如果显存不足,先把 growth_rate 降到 12,DenseBlock 的层数从 6 降到 4;把输入切片从 256x256 缩到 224x224;开启混合精度训练,gradient checkpointing 也能缓解但会拖慢速度。真的想跑 3D 版 DenseUnet,建议从 patch 尺寸 64x64x32 开始验证显存,而不是直接上全尺寸体数据,否则连前向传播都过不去。
6. 把结果做到能交差:推理、TTA 与后处理细节
推理阶段不要直接对每一张切片独立预测然后拼回三维体数据,切片之间缺少一致性,器官边缘会在 z 轴方向上出现明显的台阶感。我通常会用滑窗加重叠的方式:沿切片方向每次移动半窗,对重叠区域的多份概率取平均,这样既平滑又不需要额外训练。TTA 也很有用,水平翻转后把两个方向的预测概率平均一下,Dice 通常能稳定提升 0.5 到 1 个百分点,代价是推理时间翻倍,线上服务要谨慎开启。
小器官的后处理放在类别概率上做,不要直接在类别 ID 上做。先取每个像素最大概率的类别,然后对每一类器官做连通域分析,把体积极小的孤立簇删掉。比如胰腺这种器官应该是连续结构,如果预测出十几个小孤岛,基本是噪声,直接删除并把像素改成周围背景类的概率最高值。形态学开运算可以弥合细小断裂,但操作的核只能取一个体素大小,太大反而会把胰腺这类薄结构直接抹掉。
我在实际项目里的习惯是:训练阶段不追求某个数据集上的排行榜分数,而是把所有器官类别单独打印,每一轮看着小器官的数值变化决定调参方向。医学影像分割的落地,真正重要的不是模型结构多新颖,而是数据有没有按患者隔离、标签类别有没有对齐、窗宽窗位有没有统一。这三个坑填平之后,再谈 DenseUnet 还是 Vision Transformer 才有意义。希望这篇实战笔记能帮你在腹部多脏器分割项目上少走几趟冤枉路。
本文还有配套的精品资源,点击获取