Skip to content

Repository files navigation

VoxelMorph 无监督医学图像配准复现与形变可信化探索

本仓库是一个围绕 VoxelMorph 的 PyTorch 复现与研究探索工程。当前重点不是声称提出新算法,而是把医学图像非刚性配准的完整链路跑通,并围绕 图像相似度、解剖结构重叠、形变场拓扑合理性和结果可信性 做系统化实验记录。

项目主题:

VoxelMorph 无监督医学图像配准 PyTorch 复现与改进探究工程:原生基线 OASIS 脑 MRI 数据集复现、2D 脑部成对配准训练链路、正则强度消融、SVF integration 初步验证、Dice / Jacobian folding / smoothness 等形变量化评估,以及面向可信配准的动态模型选择、ICE 可逆性诊断、Mamba-SSM 主干替换等后续研究设想。

当前状态

已完成:

  • synthetic 2D smoke test:验证 moving/fixed -> network -> displacement -> warp -> loss -> optimizer 最小训练闭环。
  • Neurite-OASIS 2D 复现:414 张脑 MRI 2D slice,scan-to-scan 无监督配准。
  • OASIS 2D 四组训练:lambda_smooth = 0.001 / 0.01 / 0.1integration_steps = 0 / 7
  • SVF integration 对比:直接位移场与 integration_steps=7 的 stationary velocity field 积分路线。
  • 形变可信性评估:Jacobian determinant、folding rate、Jacobian 分布、形变平滑度。
  • 补充评估:30 pair multi-pair evaluation、强度扰动鲁棒性、seg4/seg24 Dice 结构重叠评估。
  • AutoDL 环境配置入口与实验结果整理;实例备份、临时日志和本地运行记录不纳入版本控制。

仍是后续设想:

  • Mamba/SSM 主干替换消融。已有 MambaMorph / TransMorph 等相关工作,因此这里只作为待验证假设,不作为已完成创新。
  • ICE self-inverse / swap-pair consistency 的正式脚本化评估。
  • 完整 3D volume benchmark 和多中心临床数据验证。
  • 学习式快速初值 + 传统优化精修的工程闭环。

项目结构

voxelmorph/
  nn/
    models.py              VxmPairwise 等主模型
    modules.py             U-Net, SpatialTransformer, IntegrateVelocityField
    losses.py              MSE, NCC, Grad smoothness
    functional.py          PyTorch 后端空间变换函数
  py/
    utils.py               NIfTI/npz 读写、Dice、Jacobian 等工具
    generators.py          数据生成器参考实现

experiments/
  README.md                实验脚本与结果索引
  smoke_vxm2d.py           synthetic 2D 最小闭环实验
  train_oasis2d.py         Neurite-OASIS 2D 无监督训练脚本
  eval_oasis2d_multipair.py    30 pair 批量评估
  eval_oasis2d_robustness.py   强度扰动鲁棒性模拟
  eval_oasis2d_dice.py         seg4/seg24 Dice 评估
  eval_oasis2d_jacobian_distribution.py Jacobian 分布统计
  train_oasis3d_prelim.py      小规模 3D preliminary 脚本
  results/oasis2d/         轻量公开复现结果
  results/oasis2d_supplement/          结果图、补充评估表格和精选可视化

scripts/
  train.py                 原 VoxelMorph 风格训练入口参考
  register.py              原 VoxelMorph 风格推理入口参考

tests/
  test_*.py                模型、模块、functional、neurite 集成测试

docs/
  docs/                    MkDocs 源文档

本地数据、训练输出、模型权重和 AutoDL 完整备份不进入 Git:

data/
outputs/
autodl_results/
autodl_backups/
*.pt
*.nii.gz
*.tar.gz

方法链路

无监督配准输入一对图像:

moving image M, fixed image F

核心前向链路:

M, F
-> concat([M, F])
-> U-Net / registration backbone
-> dense displacement field u
   或 stationary velocity field v
-> optional SVF integration: u = exp(v)
-> SpatialTransformer(M, u)
-> warped moving image

当前 OASIS 2D 训练目标:

loss = MSE(warped_moving, fixed) + lambda_smooth * smoothness(displacement)

指标分工:

维度 指标 回答的问题
图像相似度 MSE / PSNR / NCC warped moving 和 fixed 在灰度/结构上是否更接近
解剖结构重叠 Dice 如果有 segmentation label,脑区标签是否更重合
形变合理性 Jacobian determinant / folding rate / smoothness 为了让图像变像,模型是否产生局部折叠或过激形变
后续可信性 ICE / 扰动一致性 / 模型分歧 正反向形变是否一致,结果是否稳定

快速开始

1. 安装

建议使用 Python 3.10+ 和已有 GPU 版 PyTorch 的环境。AutoDL 镜像通常已带 PyTorch,不建议随意覆盖。

pip install nibabel matplotlib scipy scikit-image tqdm h5py
pip install "git+https://github.com/adalca/neurite.git@dev"
pip install -e . --no-deps

AutoDL 环境可参考:

setup_autodl.sh

公开仓库只保留 setup_autodl.sh 作为环境配置入口;实例租用、备份恢复和临时运行记录不纳入版本控制。

2. Synthetic smoke test

python experiments/smoke_vxm2d.py \
  --steps 120 \
  --out-dir outputs/smoke_vxm2d

3. 下载 Neurite-OASIS 2D

mkdir -p data
cd data
wget -c https://surfer.nmr.mgh.harvard.edu/ftp/data/neurite/data/neurite-oasis.2d.v1.0.tar
tar -xf neurite-oasis.2d.v1.0.tar
cd ..
find data -name slice_norm.nii.gz | wc -l

预期图像数量为 414。若需要 Dice 评估,还需要确认标签文件存在:

find data -name "slice_seg*.nii.gz" | head

4. OASIS 2D 四组训练

# weak regularization
python experiments/train_oasis2d.py \
  --data-root data \
  --steps 1000 \
  --batch-size 8 \
  --image-size 128 \
  --lambda-smooth 0.001 \
  --integration-steps 0 \
  --out-dir outputs/oasis2d_lam0001

# baseline
python experiments/train_oasis2d.py \
  --data-root data \
  --steps 1000 \
  --batch-size 8 \
  --image-size 128 \
  --lambda-smooth 0.01 \
  --integration-steps 0 \
  --out-dir outputs/oasis2d_baseline_lam001

# strong regularization
python experiments/train_oasis2d.py \
  --data-root data \
  --steps 1000 \
  --batch-size 8 \
  --image-size 128 \
  --lambda-smooth 0.1 \
  --integration-steps 0 \
  --out-dir outputs/oasis2d_lam01

# SVF integration
python experiments/train_oasis2d.py \
  --data-root data \
  --steps 1000 \
  --batch-size 8 \
  --image-size 128 \
  --lambda-smooth 0.01 \
  --integration-steps 7 \
  --out-dir outputs/oasis2d_int7_lam001

5. 补充评估

python experiments/eval_oasis2d_multipair.py \
  --data-root data \
  --runs-root outputs \
  --num-pairs 30 \
  --image-size 128 \
  --out-dir outputs/oasis2d_multipair_eval

python experiments/eval_oasis2d_robustness.py \
  --data-root data \
  --runs-root outputs \
  --run-dirs oasis2d_baseline_lam001 oasis2d_int7_lam001 \
  --num-pairs 20 \
  --image-size 128 \
  --out-dir outputs/oasis2d_robustness_eval

python experiments/eval_oasis2d_jacobian_distribution.py \
  --data-root data \
  --runs-root outputs \
  --num-pairs 30 \
  --image-size 128 \
  --out-dir outputs/oasis2d_jacobian_hist

python experiments/eval_oasis2d_dice.py \
  --data-root data \
  --runs-root outputs \
  --seg-name slice_seg4.nii.gz \
  --num-pairs 30 \
  --image-size 128 \
  --out-dir outputs/oasis2d_dice_eval_seg4

公开结果

轻量复现结果在:

experiments/results/oasis2d/

补充评估和结果图表在:

experiments/results/oasis2d_supplement/

Single-pair 四组结果

固定 evaluation pair:

data/OASIS_OAS1_0001_MR1/slice_norm.nii.gz
data/OASIS_OAS1_0002_MR1/slice_norm.nii.gz
Run MSE after PSNR after NCC after Jac min Folding rate
lambda=0.001, int=0 0.00544 22.65 0.9493 -12.1971 0.062561
lambda=0.01, int=0 0.00676 21.70 0.9399 -1.7743 0.024719
lambda=0.1, int=0 0.00896 20.48 0.9241 0.0298 0.000000
lambda=0.01, int=7 0.00468 23.29 0.9561 0.0355 0.000000

OASIS 2D SVF panel

Multi-pair 批量评估

30 对随机 OASIS 2D moving/fixed pair:

Run MSE after NCC after PSNR after Folding rate
lambda=0.001, int=0 0.01390 ± 0.01234 0.9326 ± 0.0258 20.12 ± 3.69 0.065788 ± 0.026025
lambda=0.01, int=0 0.01465 ± 0.01237 0.9316 ± 0.0224 19.77 ± 3.55 0.024819 ± 0.016134
lambda=0.1, int=0 0.01750 ± 0.01355 0.9193 ± 0.0219 18.79 ± 3.30 0.000694 ± 0.001493
lambda=0.01, int=7 0.01076 ± 0.01065 0.9476 ± 0.0231 21.45 ± 3.84 0.000008 ± 0.000021

Multi-pair evaluation

代表性真实影像 pair:

Multi-pair representative panels

强度扰动鲁棒性

该实验对 moving image 施加 identity、noise、gamma、invert、blur 等强度扰动,用来观察 MSE-based similarity 在 contrast shift 下的边界。它是 cross-contrast-like simulation,不是真实跨模态实验。

Perturbation baseline NCC after int7 NCC after baseline folding int7 folding
identity 0.9343 0.9531 0.01918 0.00004
noise 0.9095 0.9315 0.01832 0.00002
gamma_dark 0.9034 0.9182 0.02449 0.00006
gamma_bright 0.9239 0.9363 0.03307 0.00009
invert -0.0676 0.1651 0.20505 0.01964
blur 0.9400 0.9552 0.01382 0.00001

Robustness perturbation summary

Identity vs invert visual comparison

Jacobian 分布

Jacobian determinant 不只看 folding rate 这一个点,也可以看分布宽度、低分位数和高分位数。弱正则的分布更宽,负 Jacobian 更多;强正则更保守;SVF-int7 在当前 2D 设置下 folding rate 接近 0。

Jacobian distribution summary

Jacobian Folding 风险视角

Folding rate 可以作为形变场拓扑风险的直接信号:如果只看 NCC/MSE,弱正则模型可能显得更“贴合”图像,但同时产生更高比例的负 Jacobian 区域;如果加入 folding rate 作为横轴,就能看到图像相似度和形变合理性之间的 trade-off。当前结果中 integration_steps=7 在保持较高 NCC 的同时 folding rate 接近 0,但这仍是 OASIS 2D slice-level 证据,不能替代完整 3D diffeomorphic validation。

Jacobian folding trade-off

Jacobian heatmap 则提供空间层面的可视化:红/蓝区域显示局部面积扩张或压缩更明显的位置;当局部 Jacobian 低于 0 时,意味着该区域发生 folding。它适合和 folding rate 一起读:前者看“哪里变形剧烈”,后者看“折叠风险占比有多大”。

Jacobian heatmap montage

Dice 结构重叠评估

Dice 依赖 OASIS 2D 标签文件 slice_seg4.nii.gzslice_seg24.nii.gz,只作为离线结构级评估,不参与无监督训练。

Run seg4 Dice after seg24 Dice after Folding rate
lambda=0.001, int=0 0.6518 ± 0.0905 0.5163 ± 0.1071 0.056610 ± 0.022501
lambda=0.01, int=0 0.6658 ± 0.0832 0.5422 ± 0.0981 0.020046 ± 0.016916
lambda=0.1, int=0 0.6818 ± 0.0676 0.5765 ± 0.0907 0.000452 ± 0.001100
lambda=0.01, int=7 0.6770 ± 0.1110 0.5528 ± 0.1218 0.000006 ± 0.000019

Dice controlled seg24 panel

Dice controlled seg4 panel

结果解读

  • lambda_smooth 是控制图像相似度与形变平滑性的关键旋钮。弱正则更容易追求像素贴合,但 folding 风险更高;强正则更保守,通常能降低 folding 风险。
  • integration_steps=7 的 SVF 路线在当前 OASIS 2D 设置下取得更低 folding rate,并在 multi-pair 图像相似度上表现较好。这里不能直接宣称“完全微分同胚安全”,仍需 Jacobian、ICE 和 3D 评估验证。
  • Dice 结果显示,结构级重叠和 MSE/NCC 不总是同一个排序。医学配准不能只看图像是否变像,还需要同时检查标签重叠和形变拓扑合理性。
  • 强度扰动实验是 cross-contrast-like simulation,不是真实 CT-MRI 跨模态配准。它主要说明 MSE-based similarity 在强对比度变化下存在边界。

证据边界

  • 当前主要结果是 OASIS 2D slice-level evaluation,不等价于完整 3D volume benchmark。
  • 公开表格中的 single-pair 指标用于复现链路检查和消融示例,不应当作为泛化结论。
  • Multi-pair 和 Dice 评估增强了证据,但样本规模、标签粒度和数据域仍有限。
  • 无标签可信评分、动态模型选择和多证据风险控制是后续研究方向;当前 README 只保留方向性描述,不把未验证方案包装成已完成方法。
  • Mamba/SSM、ICE 路由、连续 lambda 控制和传统优化精修目前是后续研究计划,不是已经完成的性能结论。
  • 原始数据、训练权重、完整 AutoDL 备份和临时分析文档不上传到 GitHub。

测试

测试依赖:

pip install pytest
python -m pytest tests

若只做语法检查:

python -m py_compile \
  experiments/smoke_vxm2d.py \
  experiments/train_oasis2d.py \
  experiments/eval_oasis2d_multipair.py \
  experiments/eval_oasis2d_robustness.py \
  experiments/eval_oasis2d_dice.py \
  experiments/eval_oasis2d_jacobian_distribution.py

引用

本项目基于 VoxelMorph 相关论文和开源实现进行学习与复现:

  • Balakrishnan et al., VoxelMorph: A Learning Framework for Deformable Medical Image Registration, IEEE TMI, 2019.
  • Dalca et al., Unsupervised Learning for Fast Probabilistic Diffeomorphic Registration, MICCAI, 2018.
  • Hoopes et al., HyperMorph: Amortized Hyperparameter Learning for Image Registration, IPMI, 2021.
  • VoxelMorph GitHub: https://github.com/voxelmorph/voxelmorph
  • Neurite-OASIS dataset: https://github.com/adalca/medical-datasets

About

VoxelMorph 无监督医学图像配准 PyTorch 复现与改进探究工程:原生基线 OASIS 脑 MRI 数据集完整复现、Mamba-SSM 主干替换消融对照、微分同胚 SVF 积分 (积分步数 7) 验证、2D 脑部成对配准训练链路;配套形变场量化评测指标(Dice 系数、Jacobian 折叠率、形变平滑损失),留存非刚性配准学习笔记、实验日志与代码探索记录,欢迎讨论交流。

Topics

Resources

Stars

2 stars

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages