TransMorph: Transformer for Medical Image Registration

keywords: Vision Transformer, Swin Transformer, convolutional neural networks, image registration

This is a PyTorch implementation of my paper:

Chen, Junyu, et al. "TransMorph: Transformer for Medical Image Registration. " arXiv, 2021.

TransMorph

TransMorph DIR Variants:

There are four TransMorph variants: TransMorph, TransMorph-diff, TransMorph-bspl, and TransMorph-Bayes.
Training and inference scripts are in TransMorph/, and the models are contained in TransMorph/model/.

1. TransMorph: A hybrid Transformer-ConvNet network for image registration.
2. TransMorph-diff: A probabilistic TransMorph that ensures a diffeomorphism.
3. TransMorph-bspl: A B-spline TransMorph that ensures a diffeomorphism.
4. TransMorph-Bayes: A Bayesian uncerntainty TransMorph that produces registration uncertainty estimate.

TransMorph Affine Model:

The scripts for TransMorph affine model are in TransMorph_affine/ folder.

train_xxx.py and infer_xxx.py are the training and inference scripts for TransMorph models.

Baseline Models:

We compared TransMorph with eight baseline registration methods + four Transformer architectures.
Baseline registration methods:

1. SyN (ATNsPy)
2. NiftyReg
3. LDDMM
4. deedsBCV
5. VoxelMorph-1 & -2
6. CycleMorph
7. MIDIR

Baseline Transformer architectures:

1. PVT
2. nnFormer
3. CoTr
4. ViT-V-Net

Training and inference scripts for the baseline models will be available in the near future!

Dataset:

Due to restrictions, we cannot distribute our brain MRI data. However, several brain MRI datasets are publicly available online: IXI, ADNI, OASIS, ABIDE, etc. Note that those datasets may not contain labels (segmentation). To generate labels, you can use FreeSurfer, which is an open-source software for normalizing brain MRI images. Here are some useful commands in FreeSurfer: Brain MRI preprocessing and subcortical segmentation using FreeSurfer.

Citation:

If you find this code is useful in your research, please consider to cite:

@misc{chen2021transmorph,
title={TransMorph: Transformer for Medical Image Registration}, 
author={Junyu Chen and Yufan He and Eric C. Frey and Ye Li and Yong Du},
year={2021},
eprint={2111.10480},
archivePrefix={arXiv},
primaryClass={eess.IV}
}

TransMorph Architecture:

Example Results:

Qualitative comparisons:

Uncertainty Estimate by TransMorph-Bayes:

Quantitative Results:

Inter-patient Brain MRI:

XCAT-to-CT:

Reference:

Swin Transformer
easyreg
MIDIR
VoxelMorph

About Me