wplll

这是我(wpl/peilwang)与 Arunodhay/arunodhayan 在 Vesuvius Challenge - Surface Detection 中获得第三名的方案中文翻译。

首先感谢主办方举办这场精彩的比赛。我们很荣幸在这里分享方案:它以深度定制的 nnU-Net v2 为核心, 结合训练与推理的非对称 patch、两阶段级联预测、自定义测试时增强,以及基于 3D Hessian 矩阵的 ridge detection 后处理。

1. 整体架构流程

我们构建了一套高效、稳健的两阶段级联 3D 图像分割流程:

3D volume
  -> Stage 1:3D Fullres 初始预测与多模型概率图加权融合
  -> Stage 2:3D Cascade Fullres,接收 Stage 1 预测作为空间先验
  -> 后处理:Distance Transform + 3D Hessian 的拓扑优化与细微结构提取

2. 数据集策略与处理

全量数据训练。 为尽可能贴合真实数据分布,我们放弃传统 K-Fold 交叉验证,即不保留验证集; 最终提交模型直接使用全部训练数据训练,以获得更稳健的泛化能力。

伪标签实验。 比赛期间,我们曾使用外部/无标签数据生成伪标签以扩充训练集。实验结果未带来显著 提升,因此最终版本没有采用伪标签策略。

数据增强。 在 nnU-Net 默认增强基础上,我们针对目标的拓扑特点,强制启用所有轴的翻转 (Mirror Axes: 0、1、2),并将旋转增强概率设为较高的 0.4–0.8

3. 模型架构与训练策略

我们对默认 nnU-Net Trainer 做了深度定制,核心改进如下。

非对称 patch 策略。 我们采用 M-size 网络,并使用在 CZII 比赛中验证有效的尺寸解耦:训练 patch 设为 128,推理 patch 增至 192。这既维持训练时较高的 batch size 与迭代效率,也让推理获得更大的 感受野,从而改善分割结果的全局一致性。

拓扑感知损失。 对于连通性要求很强的结构,经典 Dice + CE 往往难以完整保留细长或网状结构。 因此,我们在默认损失上加入 clDice Loss;它在骨架层面对管状/网状结构的拓扑连通性进行专门优化。

优化器。 我们放弃传统学习率调度器,改用 RAdamScheduleFree。实践表明,它能让收敛更平稳, 也显著提高训练效率。

级联训练路径。 Stage 2 Cascade 严格遵循“由粗到细”的范式:先训练 Lowres 模型,再在其基础上 训练 Cascade 模型,使网络正确学习如何把上一阶段的粗预测作为可靠的空间先验。

4. 推理与模型融合(TTA)

为兼顾预测精度和推理速度,我们重写了 predict_from_raw_data.py 的预测逻辑。

非对称推理与阶段融合。 推理 patch 使用 192。在两阶段流程中,我们将 Fullres 模型的输出直接 作为 Cascade 模型的特征先验输入,以获得最终预测;这套组合能够捕获最丰富的图像细节。

自定义 TTA。 除常规镜像翻转外,我们通过 Affine Grid Sample 加入平面内 ±15° 旋转 TTA。每个 patch 的最终预测融合原始、翻转与 ±15° 旋转版本的输出。训练中使用 ±30° 随机旋转,实验表明 将 TTA 旋转角设为 ±15° 最优;训练中进一步增大旋转范围反而会降低性能。

并发提取与多模型融合。 我们编写专门的并行推理脚本,在多 GPU 环境中将多个进程绑定到指定显卡, 并行提取 .npz 概率图;随后对不同迭代阶段或特定设置的 checkpoint 做等权或加权融合。

5. 拓扑层面的后处理

这是本方案取得明显分数提升的关键部分。考虑到目标的强空间连通性,我们设计了以下后处理流程:

  1. Inverse EDT Dilation: 利用基于欧氏距离变换(EDT)的反向映射,对二值结果进行非线性膨胀。
  2. 3D Hessian Ridge Detection: 将目标视为空间中的“脊线”。对预测体做 Gaussian 平滑后计算 3D Hessian 矩阵,并利用其特征值构建特征空间,提取具有管状特征的结构,大幅过滤无结构背景噪声。
  3. 拓扑与一致性剪枝:
    • Z 轴一致性剪枝: 当前 Z 切片为正、上下相邻切片均为负的体素会被判为假阳性并移除,这一步 对清除跨切片离群噪声至关重要。
    • 对切片进行 2D 形态学 closing 平滑边界;同时在整体体积的 XY 平面做各向异性 3D closing,修复 断裂部分。
    • 最后移除体积过小的离散 3D 噪声。

特别致谢:后处理流程中的许多灵感与核心代码逻辑参考并受益于主办方在讨论区分享的讨论和代码。 感谢主办方和社区的无私分享。

6. 致谢

再次感谢主办方提供这项富有挑战的平台,也感谢队友 @arunodhayan 在模型训练与探索过程中的投入。

最佳 Private LB 提交: 我们的最高 private leaderboard 分数来自充分训练的 Cascade 配置: 先训练 lowres 网络 8,000 epochs,再在此基础上训练 cascade 网络 8,000 epochs。

项目链接