本仓库是一个围绕 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.1与integration_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 / 扰动一致性 / 模型分歧 | 正反向形变是否一致,结果是否稳定 |
建议使用 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-depsAutoDL 环境可参考:
setup_autodl.sh
公开仓库只保留 setup_autodl.sh 作为环境配置入口;实例租用、备份恢复和临时运行记录不纳入版本控制。
python experiments/smoke_vxm2d.py \
--steps 120 \
--out-dir outputs/smoke_vxm2dmkdir -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# 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_lam001python 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/
固定 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 |
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 |
代表性真实影像 pair:
该实验对 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 |
Jacobian determinant 不只看 folding rate 这一个点,也可以看分布宽度、低分位数和高分位数。弱正则的分布更宽,负 Jacobian 更多;强正则更保守;SVF-int7 在当前 2D 设置下 folding rate 接近 0。
Folding rate 可以作为形变场拓扑风险的直接信号:如果只看 NCC/MSE,弱正则模型可能显得更“贴合”图像,但同时产生更高比例的负 Jacobian 区域;如果加入 folding rate 作为横轴,就能看到图像相似度和形变合理性之间的 trade-off。当前结果中 integration_steps=7 在保持较高 NCC 的同时 folding rate 接近 0,但这仍是 OASIS 2D slice-level 证据,不能替代完整 3D diffeomorphic validation。
Jacobian heatmap 则提供空间层面的可视化:红/蓝区域显示局部面积扩张或压缩更明显的位置;当局部 Jacobian 低于 0 时,意味着该区域发生 folding。它适合和 folding rate 一起读:前者看“哪里变形剧烈”,后者看“折叠风险占比有多大”。
Dice 依赖 OASIS 2D 标签文件 slice_seg4.nii.gz 和 slice_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 |
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 pytestpython -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









