nnU-Net + HD-BET
脑 MRI 分割实战记录

本报告完整记录在一台无 GPU的 Linux 工作站上， 跑通两套医学图像分割开源工具的全过程：包括所有命令、超参数、 数据流转、踩坑与解决方案，以及对推理结果的定量评估。

§0摘要

本次实验完成两个独立的医学图像分割任务的端到端推理：

1. Track A · nnU-Net 海马体分割：基于 Medical Segmentation Decathlon (MSD) Task04 数据集， 使用 DKFZ 团队 2020 年公开的预训练权重，在 CPU 上对 6 例（3 训练集 + 3 测试集）做了真分割推理， 单例平均耗时约 0.95 秒，训练集子集上 Dice 系数全部 ≥ 0.948。
2. Track B · HD-BET 颅骨剥离：使用 HD-BET 自带预训练模型， 对 ICBM152 标准脑模板 (197×233×189 voxel) 做颅骨剥离推理， 输出脑实质体积占总体素 21.8%。

nnU-Net 训练步骤本次未执行（需要 GPU），但完成了所有训练前准备工作（数据指纹生成 + auto plan）。

§1技术栈

§2数据集 · MSD Task04 Hippocampus

Medical Segmentation Decathlon 第 04 号任务，由 Vanderbilt University Medical Center 提供， 目标是对 T1 加权脑 MRI 中的前海马（Anterior, label=1）和后海马 （Posterior, label=2）进行体素级分割。

体素强度统计（前景区域，所有 260 例）

这个分布充分说明 nnU-Net 强制要求 Z-Score 归一化是合理的： 原始 MRI 强度跨越 6 个数量级，不归一化网络根本没法收敛。

A.1 · 数据准备（下载 + 转 v2 格式）

原始 MSD 数据集使用旧版 schema（tensorImageSize、modality 等字段名）， nnU-Net v2 不能直接消费。需要先转换。

# 1. 从 MSD AWS bucket 下载
$ wget https://msd-for-monai.s3.amazonaws.com/Task04_Hippocampus.tar
$ tar -xf Task04_Hippocampus.tar -C downloads/

# 2. 转换到 nnU-Net v2 schema
$ nnUNetv2_convert_MSD_dataset \
       -i downloads/Task04_Hippocampus \
       -overwrite_id 4

转换脚本做了三件事：

• 每个图像文件加 _0000 后缀，标识"第 0 个模态通道"
• 重写 dataset.json： modality → channel_names， labels 列表 → 字典格式
• 顶层目录改名：Task04_Hippocampus → Dataset004_Hippocampus

新 dataset.json（节选）

{
  "channel_names": { "0": "MRI" },
  "labels": {
    "background": 0,
    "Anterior":  1,
    "Posterior": 2
  },
  "numTraining": 260,
  "file_ending": ".nii.gz"
}

A.2 · 数据指纹（dataset_fingerprint.json）

nnUNetv2_plan_and_preprocess 第一步会扫描全部 260 例， 生成数据指纹 — 一个描述数据集统计特征的 JSON， 作为后续 plan 生成的输入。

$ nnUNetv2_plan_and_preprocess -d 4 --verify_dataset_integrity

# 内部步骤：
#   1) verify_dataset_integrity   - 校验所有 image/label 一一对应
#   2) extract_fingerprint        - 体素统计 + 形状统计
#   3) plan_experiments           - 根据指纹推导网络架构
#   4) preprocess                 - 重采样 + 归一化 + 裁剪 + 压缩存盘

形状分布（260 例）

数据集体积非常小（中位 36×50×35 ≈ 6.3 万 voxel）， 这也是为什么本次能在 CPU 上做实时推理。

A.3 · 自动训练计划（nnUNetPlans.json）

这是 nnU-Net 的核心创新：根据数据指纹自动推导出网络架构、patch size、batch size 等所有超参数， 全程只用规则（不依赖梯度搜索），花费 CPU 时间约几秒。本次同时生成了 2D 和 3D 两套配置。

网络架构 (3d_fullres)

network_class_name: dynamic_network_architectures.architectures.unet.PlainConvUNet
n_stages:            4          # 编码器 4 级下采样
features_per_stage: [32, 64, 128, 256]   # 翻倍递增
conv_op:             torch.nn.Conv3d
kernel_sizes:        all 3×3×3
strides:             stage 0 = 1×1×1, stages 1-3 = 2×2×2
n_conv_per_stage:    2 (encoder) / 2 (decoder)
norm_op:             InstanceNorm3d
dropout_op:          None
activation:          LeakyReLU(slope=0.01)

A.4 · 训练（本次跳过）

完整训练命令在脚本里写好了，但本次未执行。原因：5 折交叉验证 × 1000 epoch × 250 iter/epoch 在 CPU 上预估需要数周甚至更久。Hippocampus 数据集小，估计 GPU 上单 fold 1-2 小时。

# 训练命令（本次未执行）
$ for fold in 0 1 2 3 4; do
      nnUNetv2_train 4 2d $fold --npz
  done

# 替代方案 → §A.5

A.5 · 预训练权重（DKFZ Zenodo）

nnU-Net 团队 2020 年发布论文时，把 MSD 全部 10 个任务训练好的权重打包上传到 Zenodo （record 4003545）， 许可证 CC-BY 4.0。本次直接使用 Task004_Hippocampus.zip（271 MB）。

$ wget https://zenodo.org/records/4003545/files/Task004_Hippocampus.zip

# 解压后包含：
#   - 5 fold 交叉验证训练好的 model_final_checkpoint.model
#   - 同样 5 fold 的 2d 模型（双配置 ensemble）
#   - postprocessing.json（含官方原始 cross-val Dice）

$ nnUNet_install_pretrained_model_from_zip Task004_Hippocampus.zip

权重 zip 文件结构

Task004_Hippocampus.zip
├── 3d_fullres/Task004_Hippocampus/nnUNetTrainerV2__nnUNetPlansv2.1/
│   ├── plans.pkl                    # v1 时代的 plans (321 KB)
│   ├── postprocessing.json          # 官方 cross-val Dice
│   ├── fold_0/
│   │   ├── model_final_checkpoint.model       # 45 MB · 网络权重
│   │   ├── model_final_checkpoint.model.pkl   # trainer state
│   │   ├── progress.png                       # 训练曲线
│   │   └── network_architecture.pdf
│   ├── fold_1/ ... fold_4/         # 5 折交叉验证
├── 2d/Task004_Hippocampus/...      # 2D 配置同结构
└── ensembles/...                   # 2D + 3D ensemble 配置

A.6 · CPU 推理

分两批跑：训练集 3 例（用于评估，因为有 GT 可对比 Dice）+ 测试集 3 例（演示真实推理）。

# 关键环境变量（v1 风格命名）
export nnUNet_raw_data_base=/tmp/v1_workspace/nnUNet_raw_data_base
export nnUNet_preprocessed=/tmp/v1_workspace/nnUNet_preprocessed
export RESULTS_FOLDER=/tmp/v1_workspace/RESULTS_FOLDER

# 推理（fold 0 单模型，关 TTA 提速）
$ nnUNet_predict \
       -i /tmp/v1_inference/input \
       -o /tmp/v1_inference/output \
       -t Task004_Hippocampus \
       -m 3d_fullres \
       -f 0 \
       --disable_tta

推理时间（CPU · 单 fold · disable_tta）

prediction_time.txt 由 nnU-Net 自动写入，是 nnUNet_predict 内部计时器记录的纯网络前向传播时间， 不包括磁盘 I/O。

A.7 · Dice 评估

Dice 系数（Dice-Sørensen Coefficient）衡量两个集合的体素重叠程度： Dice = 2|A ∩ B| / (|A| + |B|)， 值域 [0, 1]，1.0 = 完美重叠。本次用 numpy 直接算每个 label 的 Dice：

def dice(pred, gt, label):
    p = pred == label
    g = gt   == label
    inter = (p & g).sum()
    return 2 * inter / (p.sum() + g.sum() + 1e-9)

训练集 Dice 结果（3 例）

原作者发布的真实评估指标

从权重 zip 内的 postprocessing.json 解析出官方 5 折交叉验证 Dice：

原作者公布的 cross-val Dice ≈ 0.89/0.88 才是模型对未见数据的真实表现， 这也是 MSD challenge leaderboard 上的水平。本次跑的 0.96 只是流水线验证， 不应理解为模型性能。

测试集推理（无 GT，仅展示）

另外对测试集的 case 002/005/009 也做了推理，输出已写入磁盘 （/tmp/v1_inference/output_test/hippocampus_*.nii.gz）。 由于 MSD 测试集不公开标注，无法计算 Dice，但分割结果可以用 ITK-SNAP 目视检查。

A.8 · 踩坑记录

坑 1 · v1 权重不能直接喂 v2

最初想用 v2 推理 v1 权重，提示找不到 trainer 类。原因：v1 训练器是 nnUNetTrainerV2（v2 时代不存在）， 且目录结构、plans 文件名都不一样。

解决：装一个独立的 v1 venv 专门做推理，预处理仍用 v2。

坑 2 · PyTorch 2.6 weights_only 默认值变更

_pickle.UnpicklingError: Weights only load failed.
  Unsupported global: GLOBAL numpy.core.multiarray.scalar was not an allowed global
  by default.

PyTorch 2.6 把 torch.load() 的 weights_only 参数默认改成了 True（之前是 False）， 老 checkpoint 里的 numpy scalar 对象在白名单外被拒绝加载。

解决：给 nnU-Net v1 的 model_restore.py 打补丁：

# site-packages/nnunet/training/model_restore.py · line 147
- all_params = [torch.load(i, map_location=...) for i in ...]
+ all_params = [torch.load(i, map_location=..., weights_only=False) for i in ...]

坑 3 · numpy 2.x 不兼容

PyTorch 当前 PyPI wheel 是用 numpy 1.x 编译的。如果 pip 自动装了 numpy 2.x， 导入 torch 会报 _ARRAY_API not found。必须 pip install "numpy<2"。

坑 4 · HD-BET 参数风格

HD-BET CLI 的设备参数是 -device cpu（单短横线），不是 --device cpu。 这是 argparse 的可选位置参数风格，跟一般 Python 工具的双短横线惯例不一样。

B.1 · 工具背景

HD-BET 是 DKFZ 团队（同 nnU-Net 作者）发布的脑提取工具，论文： Isensee et al., HBM 2019。 网络架构基于 nnU-Net 的 3D U-Net，训练在 1568 例多中心 MRI 上。

解决的问题：原始脑 MRI 包含头骨、皮肤、眼球、脖子等非脑组织， 对后续分析（分割海马体、检测肿瘤、配准到标准空间等）都是噪声。 HD-BET 把这些剥掉，只留 brain mask 内的组织。

临床意义：这是几乎所有脑 MRI 流水线的第一步预处理。 传统方法（FSL BET、AFNI 3dSkullStrip）参数敏感，效果不稳定；HD-BET 是 SOTA 之一， 论文报告 Dice ≈ 0.96 vs FSL BET ≈ 0.92。

B.2 · 推理过程

$ pip install hd-bet

# 首次运行：自动从作者服务器下载权重 (~400 MB)
#   缓存到 ~/.cache/torch/hub/checkpoints/

$ hd-bet \
       -i hdbet_test/test_brain.nii.gz \
       -o hdbet_test/output/brain.nii.gz \
       -device cpu \
       --disable_tta \
       --save_bet_mask

关键参数详解

默认开启 TTA 时 HD-BET 会做 8 次镜像翻转推理然后投票，CPU 上慢 8×。 --disable_tta 提速 4×，论文报告 Dice 损失 < 0.005，CPU 推荐打开。

B.3 · 输出分析

对 ICBM152 标准脑模板（共 8,675,289 voxel）的推理结果：

ICBM152 是多人脑模板的平均，所以略大于个体大脑。21.8% 这个比例符合预期（通常脑实质占头部 20-25%）。

§4产物清单

nnunet_demo/
├── downloads/Task04_Hippocampus/          # A · MSD 原始 (56 MB)
├── nnUNet_raw/Dataset004_Hippocampus/     # A · v2 格式 (29 MB)
│   ├── imagesTr/hippocampus_*_0000.nii.gz   # 260 训练
│   ├── labelsTr/hippocampus_*.nii.gz        # 260 标注
│   ├── imagesTs/                            # 130 测试
│   └── dataset.json                         # v2 schema
├── nnUNet_preprocessed/Dataset004_Hippocampus/ # A · 71 MB
│   ├── nnUNetPlans.json                  # ◆ 自动训练计划
│   ├── dataset_fingerprint.json          # 数据指纹
│   ├── nnUNetPlans_2d/*.b2nd                # 2D 备料
│   └── gt_segmentations/                    # 原尺寸 GT
├── v1_workspace/RESULTS_FOLDER/nnUNet/3d_fullres/Task004_Hippocampus/
│   └── nnUNetTrainerV2__nnUNetPlansv2.1/
│       ├── fold_{0..4}/model_final_checkpoint.model  ◆ 271 MB · 5 折权重
│       └── plans.pkl
├── v1_inference/
│   ├── input/hippocampus_{001,003,004}_0000.nii.gz
│   ├── output/hippocampus_{001,003,004}.nii.gz     ◆ A 训练集推理结果
│   ├── input_test/hippocampus_{002,005,009}_0000.nii.gz
│   └── output_test/hippocampus_{002,005,009}.nii.gz ◆ A 测试集推理结果
├── hdbet_test/                            # B · HD-BET
│   ├── test_brain.nii.gz                    # 输入：ICBM152
│   └── output/test_brain_bet*.nii.gz       ◆ B 输出
└── outputs/
    ├── hippocampus_pred_vs_gt.png         ◆ 图 3
    ├── hdbet_result.png                   ◆ 图 4
    ├── nnunet_plan.png                      # 图 2
    └── hippocampus_sample.png               # 图 1

§5复现指南

在一台干净的 Linux/macOS 机器上复现这套流程，预计 20-30 分钟。

# 0. 环境（推荐隔离 venv）
$ python3 -m venv nnunet_v2_venv
$ source nnunet_v2_venv/bin/activate
$ pip install nnunetv2 hd-bet "numpy<2" torch torchvision

# 1. 设置环境变量（v2 风格）
$ export nnUNet_raw=$(pwd)/nnUNet_raw
$ export nnUNet_preprocessed=$(pwd)/nnUNet_preprocessed
$ export nnUNet_results=$(pwd)/nnUNet_results

# 2. Track A · 数据准备
$ wget https://msd-for-monai.s3.amazonaws.com/Task04_Hippocampus.tar
$ tar -xf Task04_Hippocampus.tar -C downloads/
$ nnUNetv2_convert_MSD_dataset -i downloads/Task04_Hippocampus -overwrite_id 4
$ nnUNetv2_plan_and_preprocess -d 4 --verify_dataset_integrity

# 3. Track A · 推理（需要切到 v1 venv）
$ python3 -m venv nnunet_v1_venv && source nnunet_v1_venv/bin/activate
$ pip install nnunet "numpy<2" torch
$ # 打补丁修 PyTorch 2.6 兼容（见 §A.8）
$ wget https://zenodo.org/records/4003545/files/Task004_Hippocampus.zip
$ nnUNet_install_pretrained_model_from_zip Task004_Hippocampus.zip
$ nnUNet_predict -i input/ -o output/ -t Task004_Hippocampus -m 3d_fullres -f 0 --disable_tta

# 4. Track B · HD-BET（任意 venv）
$ hd-bet -i brain.nii.gz -o brain_bet.nii.gz -device cpu --disable_tta --save_bet_mask

§6参考文献

1. Isensee, F., Jaeger, P. F., Kohl, S. A., Petersen, J., & Maier-Hein, K. H. (2021). nnU-Net: a self-configuring method for deep learning-based biomedical image segmentation. Nature Methods, 18(2), 203-211. DOI
2. Isensee, F. et al. (2019). Automated brain extraction of multi-sequence MRI using artificial neural networks. Human Brain Mapping, 40(17), 4952-4964. DOI
3. Antonelli, M. et al. (2022). The Medical Segmentation Decathlon. Nature Communications, 13, 4128. DOI
4. Zenodo · nnU-Net pretrained models for MSD · record 4003545
5. GitHub · MIC-DKFZ/nnUNet · github.com/MIC-DKFZ/nnUNet
6. GitHub · MIC-DKFZ/HD-BET · github.com/MIC-DKFZ/HD-BET