从 patch 热力图到 U-Net：肝脏识别进度记录

0. 背景

好久没更新博客了。

上一次还在看 nnU-Net 的论文，想着先把数据空间、预处理和输入格式这些基础问题弄明白。后来真正开始做肝脏识别，路线也从一开始的 patch 分类和粗 ROI，慢慢走到了现在的 U-Net 像素级分割。

这篇不打算写成很完整的教程，主要想记一下：

为什么换模型、换完以后做了什么，以及现在到底做到哪一步了。

先说明一下，目前实际训练的是自己写的轻量 2D U-Net baseline，并不是直接运行官方 nnU-Net。完整代码放在公开仓库里：liver-segmentation code。

1. 换模型前，旧方案卡在哪里

换 U-Net 之前，做的是一个监督式对比学习的 patch 分类器。

它会判断一个 96 × 96 的 CT patch 是 liver 还是 non-liver，然后在完整 CT 上滑窗，把每个 patch 的概率重新叠回三维空间，最后得到一张肝脏热力图。

这条路线不是完全没用。它确实能大致找到肝脏，也很适合验证“模型到底有没有学到肝脏区域”。但越往后做，问题越明显：

  • 它学的是“这个 patch 里像不像有肝脏”，不是“每一个像素是不是肝脏”；
  • 滑窗概率回填以后，边缘容易变厚、变糊；
  • 模型只看局部，容易把附近灰度相似的器官和软组织一起圈进来；
  • 阈值一改，Recall 和 Precision 就会明显拉扯；
  • 最大连通域、闭运算、矢状面过滤都能清掉一些小块，但更像是在补救结果。

旧方案最后选中的是 epoch 40、阈值 0.80。在 4 个完整验证 CT 上，平均 Recall 是 0.9504，Precision 是 0.7266，Dice 是 0.8220。

召回率已经不算低，但误检还是很多。这个时候继续给后处理叠规则，收益越来越有限。后来想明白，问题不只是模型容量不够，而是任务本身没有对齐：

想要的是一张像素级 mask，却一直在训练一个 patch 级分类器。

所以决定换成 U-Net。

2. 先把 U-Net 的最小闭环跑通

这次没有一上来就堆很复杂的结构，而是先写了一个比较小的 2D U-Net：

  • 单通道轴向 CT 输入；
  • 4 层 encoder 和 decoder；
  • skip connection；
  • GroupNorm；
  • base channels 为 16；
  • 总参数量约 194 万。

输出也从“这个 patch 是不是肝脏”，变成了“这一张切片上每个像素属于肝脏的概率”。

损失函数先用了 Soft F2。F2 会比普通 Dice 更看重漏检，比较符合现在的目标：第一阶段先尽量把肝脏完整圈住，再慢慢压低误检。

数据使用 SLIVER07 的 20 个带标注病例，按病例分成：

16 例训练
4 例验证

这里没有把切片随机打散以后再划分，而是直接按病例隔离，避免同一个人的相邻切片同时出现在训练集和验证集。

预处理阶段先把 CT window 固定在 [-100, 250]，再缓存成 NumPy 数组。最终训练集一共有 3,134 张轴向切片，验证集有 1,025 张。

正式训练前还先做了一个小 pilot：

512 张训练切片
3 epochs
1 个完整 CT 验证

这个阶段主要不是追求分数，而是检查缓存、缩放、mask、loss、反向传播和整卷推理能不能闭环。跑通以后，又做了一次阈值扫描，最后把全量训练的验证阈值定在 0.30。

3. 全量 baseline 做了什么

完整 baseline 在 CUDA 上训练了 30 epochs，batch size 为 32，输入尺寸为 256 × 256。

和之前只看 patch accuracy 不同，这次每一个 epoch 结束后，都会把 4 个验证病例的完整 CT 全部推理一遍，再统计：

  • Recall；
  • Precision；
  • Dice；
  • F2；
  • 连通域数量；
  • 最大连通域占比。

最佳模型也不再按某一批切片的表现选择，而是按整卷验证的 F2 保存。

最后 best checkpoint 落在 epoch 30。

这个过程里我觉得比较重要的一点是：切片 loss 下降不代表整卷空间结果一定更好。真正要用在三维 CT 上，就应该尽早把整卷验证放进训练流程，而不是最后才随便挑一例看看。

4. 到目前为止的结果

为了和旧方案用同一个口径比较，下面都是在同一组 4 个验证病例上取宏平均：

方法RecallPrecisionDice
SupCon patch + 滑窗0.95040.72660.8220
2D U-Net raw mask0.98370.94490.9639

换成 U-Net 后：

  • Recall 提高约 3.33 个百分点；
  • Precision 提高约 21.83 个百分点；
  • Dice 提高约 14.19 个百分点。

4 个验证病例的 Dice 分别是：

sliver07_001  0.9499
sliver07_008  0.9598
sliver07_012  0.9701
sliver07_015  0.9758

更让我意外的是，这一版 U-Net 的数字来自 raw mask，还没有套最大连通域或其他清理；旧方案的结果已经做过连通域后处理。也就是说，目前的提升主要来自训练目标终于和分割任务对上了，而不是靠后处理把图修漂亮。

推理流程现在也已经可以输出：

liver_probability.nii.gz
liver_mask_raw.nii.gz
liver_contour_raw.nii.gz
inference_summary.json

这些 NIfTI 会保留原 CT 的 spacing、origin 和 direction，可以继续放进 3D Slicer 或 ITK-SNAP 里逐层检查。

5. 现在还不能高兴得太早

虽然这次结果比旧方案好很多，但目前还是一个很早期的 baseline。

现在比较明显的问题有：

  1. 数据只有 20 例，验证也只有 4 例，还没有独立测试集和外部数据集；
  2. 目前是逐张轴向切片的 2D U-Net，没有利用相邻切片和完整三维上下文；
  3. 数据增强现在主要是亮度、对比度和噪声，还不够完整；
  4. raw mask 每例仍有 27～40 个连通块，只是最大连通域已经占预测体积的 99% 以上；
  5. 目前只统计了 Recall、Precision、Dice 和 F2，还没有 HD95、Surface Dice 这些边界指标；
  6. 现在的结果只能算内部实验，不能写成已经具备临床泛化能力。

所以这次更准确的结论不是“肝脏分割做完了”，而是：

从一个比较绕的粗定位方案，走到了一个终于能稳定学习像素级肝脏边界的 baseline。

6. 接下来的目标

后面准备按这个顺序继续：

  1. 先给 4 个验证病例统一做后处理和可视化，确认小连通块到底来自哪里；
  2. 加入更完整的空间增强，并比较 Soft F2、Dice + BCE、Focal 等损失；
  3. 做 5-fold 或至少更可靠的独立测试，补上 HD95 和 Surface Dice；
  4. 尝试 2.5D、3D U-Net，并正式跑一版官方 nnU-Net 作为对照；
  5. 增加其他肝脏数据集，重点看跨数据集泛化；
  6. 等肝脏 mask 稳定以后，把它作为 ROI，继续做肝内病灶识别。

最终目标还是没有变：

先稳定找到完整肝脏
        ↓
缩小病灶搜索范围
        ↓
再做肝内病灶识别与分割

这次最大的收获大概是，换模型并不只是把一个网络名字换成另一个网络名字。更重要的是，先确认训练目标和真正想解决的问题是不是同一件事。

现在至少终于走在同一条路上了。