ViT-V-Net

keywords: vision transformer, convolutional neural networks, image registration

This is a PyTorch implementation of my short paper:

Chen, Junyu, et al. "ViT-V-Net: Vision Transformer for Unsupervised Volumetric Medical Image Registration. " arXiv, 2021.

train.py is the training script. models.py contains ViT-V-Net model.

Pretrained ViT-V-Net: pretrained model

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.

Model Architecture:

net_arch

Vision Transformer Achitecture:

trans_arch

Example Results:

ViTVNet_res

Quantitative Results:

dice_details_

Reference:

TransUnet

ViT-pytorch

VoxelMorph

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

@misc{chen2021vitvnet,
title={ViT-V-Net: Vision Transformer for Unsupervised Volumetric Medical Image Registration}, 
author={Junyu Chen and Yufan He and Eric C. Frey and Ye Li and Yong Du},
year={2021},
eprint={2104.06468},
archivePrefix={arXiv},
primaryClass={eess.IV}
}

GitHub

https://github.com/junyuchen245/ViT-V-Net_for_3D_Image_Registration_Pytorch