This is an unofficial implementation for VoxelMorph [1]. Pytroch Lightning + Monai are used to construct the training & evaluation framework to facilitate people make their own modifications and run experiment.
Monai's CacheDataset is used to accelerate training speed by 20x.
This repo built upon the official VoxelMorph-repo, but has several differences:
- Official repo does not provide a directly-runable pytroch training script. (as the time of my writing)
- Official repo does not provide a pretrained model in pytorch format.
- This repo includes the mutual information (MI) as the loss function
- This repo uses Monai framework for image preprocessing and augmentation
- This repo uses Pytorch-Lightning framework to manage more complicated experiment, which is easier for developers and researchers who want to make modifications to network architecture, training strategy or data augmentation.
Before training, I used to organize summarize training dataset information in a json file, which will be used for construct Dataset class. An exmaple can be found here. You may notice I made up 10 pseudo cases :) just to give an example.
Note that only the nii.gz path is given, and monai will automatically read it.
See this Colab example to do training and inference with those pseudo cases.
To do the training, run
cd src/scripts
# inspect all options
python train.py --help
# do training
bash run_train.shTo do the inference using pretrained model:
# modifiy the image path
python inference.py [1] VoxelMorph: A Learning Framework for Deformable Medical Image Registration
Guha Balakrishnan, Amy Zhao, Mert R. Sabuncu, John Guttag, Adrian V. Dalca
IEEE TMI: Transactions on Medical Imaging. 2019.
eprint arXiv:1809.05231
- Write more informative doc.
- Provide pytorch pretrained model on OASIS dataset for T1w/MPRAGE registration
- Provide pytorch pretrained model on HCP dataset