VoxelMorph Head CT Deformable Registration

面向“患者仰头 CTMoving到平扫 CTFixed”的 3D 形变配准工程。项目包含 DICOM 转 NIfTI、医学图像预处理、官方 VoxelMorph 训练适配器、独立推理,以及 Streamlit 交互式结果查看界面。

核心模型使用官方仓库 voxelmorph/voxelmorph 的 PyTorch 实现:

  • 官方仓库:https://github.com/voxelmorph/voxelmorph
  • 固定提交:db73f34b910bcefcb520f7f40a1bc4a3e0b6401d
  • 核心类:voxelmorph.nn.models.VxmPairwise
  • 平滑正则:neurite.nn.modules.SpatialGradient
  • 相似度:默认使用带 epsilon 的稳定 Local NCC可通过 --ncc-impl neurite 切换到 neurite.nn.modules.NCC

工程目录

Voxelmorph_Head_CT/
├── Data/
│   ├── 患者1-平扫CT/              # Fixed 原始 DICOM
│   ├── 患者1-仰头CT/              # Moving 原始 DICOM
│   └── 患者2-平扫CT/
├── app.py                         # Streamlit Web 可视化界面
├── config.py                      # 默认路径与预处理参数
├── data_loader.py                 # DICOM 序列读取与 NIfTI 转换
├── environment.yml                # Conda CUDA 环境
├── infer.py                       # 独立推理,输出 warped image 与 DDF
├── metrics.py                     # NCC/MSE/MAE/DDF 等量化指标
├── model_and_train.py             # 官方 VoxelMorph 训练适配器
├── preprocess.py                  # 重采样、窗宽窗位、裁剪/填充
├── requirements.txt
└── outputs/
    ├── nifti/                     # DICOM 转换结果
    ├── preprocessed/              # 预处理结果
    ├── checkpoints/               # 模型权重
    └── inference/                 # 推理输出

环境

推荐使用 Conda 创建 CUDA 环境。本机如有 NVIDIA 驱动,不需要单独安装系统 CUDA toolkitpytorch-cuda=12.4 会随 Conda 环境提供运行时。

为减少 Conda 包缓存和 Qt 依赖占用环境采用两步安装Conda 只安装 Python、PyTorch 与 CUDA runtime业务依赖和官方 VoxelMorph 用 pip 安装。

conda env create -f environment.yml
conda activate voxelmorph-head-ct
python -m pip install --no-cache-dir -r requirements.txt
python -c "import torch; print(torch.__version__, torch.cuda.is_available())"

也可以直接运行项目脚本:

bash scripts/setup_env.sh

如果只做 CPU 调试,也可以使用 pip install -r requirements.txt,但 3D 训练强烈建议使用 CUDA。

1. DICOM 转 NIfTI

python data_loader.py \
  --dicom-dir "Data/患者1-平扫CT" \
  --output "outputs/nifti/patient1_fixed.nii.gz"

python data_loader.py \
  --dicom-dir "Data/患者1-仰头CT" \
  --output "outputs/nifti/patient1_moving.nii.gz"

data_loader.py 会优先按 InstanceNumber 排序,其次按 SliceLocation 排序,并保存 spacing、层厚、排序依据等元数据 JSON。

2. 预处理

python preprocess.py \
  --input "outputs/nifti/patient1_fixed.nii.gz" \
  --output "outputs/preprocessed/patient1_fixed_preprocessed.nii.gz" \
  --target-spacing 1 1 1 \
  --target-shape 256 256 352

python preprocess.py \
  --input "outputs/nifti/patient1_moving.nii.gz" \
  --output "outputs/preprocessed/patient1_moving_preprocessed.nii.gz" \
  --target-spacing 1 1 1 \
  --target-shape 256 256 352

默认窗口为 W=400, L=40,适合观察颈部软组织和气道。

3. 训练

python model_and_train.py \
  --moving "outputs/preprocessed/patient1_moving_preprocessed.nii.gz" \
  --fixed "outputs/preprocessed/patient1_fixed_preprocessed.nii.gz" \
  --checkpoint "outputs/checkpoints/vxm_head_ct_patient1.pt" \
  --epochs 80 \
  --image-loss mse \
  --nb-features 8 8 8 8 8 \
  --smooth-weight 0.01

训练脚本会调用官方 vxm.nn.models.VxmPairwise。如果显存不足,可降低 --nb-features 或改小 --target-shape 后重新预处理。 由于 CT 预处理后常有大面积零填充背景,默认 --ncc-impl local 会在 NCC 方差项中加入 epsilon避免低方差窗口导致非有限 loss。

4. 推理

python infer.py \
  --moving "outputs/preprocessed/patient1_moving_preprocessed.nii.gz" \
  --fixed "outputs/preprocessed/patient1_fixed_preprocessed.nii.gz" \
  --checkpoint "outputs/checkpoints/vxm_head_ct_patient1.pt" \
  --out-dir "outputs/inference"

推理会输出:

  • outputs/inference/moving_model_input.nii.gz
  • outputs/inference/fixed_model_input.nii.gz
  • outputs/inference/warped_moving.nii.gz
  • outputs/inference/ddf_mm.nii.gz
  • outputs/inference/metrics.json

moving_model_input.nii.gzfixed_model_input.nii.gz 是进入 VoxelMorph 前的统一网格图像。即使两套 CT 原始层数不同,推理前也会按 checkpoint 的输入尺寸和目标 spacing 完成重采样、归一化、中心裁剪/填充。

5. Web 结果展示

streamlit run app.py

网页提供:

  • 患者1专用固定路径Fixed 为 患者1-平扫CTMoving 为 患者1-仰头CT
  • “重新训练模型”和“开始推理”按钮。
  • Axial、Coronal、Sagittal 正交三视图;每个平面按行同时展示 Fixed、Moving、Warped。
  • Fixed 与 Warped 的 Alpha 融合或棋盘格对比。
  • DDF 位移强度热力图。
  • NCC、MSE、MAE、逐切片误差曲线、DDF 位移分布等量化图。

内存注意事项

  • DICOM 转换和重采样都有 --max-memory-mb 防护。
  • Web 界面对超大 NIfTI 会通过 nibabel proxy 按 stride 切片读取并下采样,只影响浏览器展示,不改变磁盘结果;侧栏可调整显示体素上限。
  • 训练阶段的主要瓶颈是 3D U-Net 显存;256x256x352 是较重的 3D 输入,建议优先使用 CUDA GPU。当前患者1默认使用较轻的 8 8 8 8 8 特征配置。
Description
No description provided
Readme 250 MiB
Languages
Python 99.3%
Shell 0.7%