cwmok / C2FViT

This is the official Pytorch implementation of "Affine Medical Image Registration with Coarse-to-Fine Vision Transformer" (CVPR 2022), written by Tony C. W. Mok and Albert C. S. Chung.

Affine Medical Image Registration with Coarse-to-Fine Vision Transformer (C2FViT)

This is the official Pytorch implementation of "Affine Medical Image Registration with Coarse-to-Fine Vision Transformer" (CVPR 2022), written by Tony C. W. Mok and Albert C. S. Chung.

Prerequisites

• Python 3.5.2+
• Pytorch 1.3.0 - 1.7.1
• NumPy
• NiBabel

This code was tested with Pytorch 1.7.1 and NVIDIA TITAN RTX GPU.

Training and testing scripts

• Train_C2FViT_pairwise.py: Train a C2FViT model in an unsupervised manner for pairwise registration (Inter-subject registration).

• Train_C2FViT_pairwise_semi.py: Train a C2FViT model in an semi-supervised manner for pairwise registration (Inter-subject registration).

• Train_C2FViT_template_matching.py: Train a C2FViT model in an unsupervised manner for brain template-matching (MNI152 space).

• Train_C2FViT_template_matching_semi.py: Train a C2FViT model in an semi-supervised manner for brain template-matching (MNI152 space).

• Test_C2FViT_template_matching.py: Register an image pair with a pretrained C2FViT model (Template-matching).

• Test_C2FViT_pairwise.py: Register an image pair with a pretrained C2FViT model (Pairwise image registration).

Inference

Template-matching (MNI152):

python Test_C2FViT_template_matching.py --modelpath {model_path} --fixed ../Data/MNI152_T1_1mm_brain_pad_RSP.nii.gz --moving {moving_img_path}

Pairwise image registration:

python Test_C2FViT_pairwise.py --modelpath {model_path} --fixed {fixed_img_path} --moving {moving_img_path}

Pre-trained model weights

Pre-trained model weights can be downloaded with the links below:

Unsupervised:

• C2FViT_affine_COM_pairwise_stagelvl3_118000.pth
• C2FViT_affine_COM_template_matching_stagelvl3_116000.pth

Semi-supervised:

• C2FViT_affine_COM_pairwise_semi_stagelvl3_95000.pth
• C2FViT_affine_COM_template_matching_semi_stagelvl3_130000.pth

Train your own model

Step 0 (optional): Download the preprocessed OASIS dataset from https://github.com/adalca/medical-datasets/blob/master/neurite-oasis.md and place it under the Data folder.

Step 1: Replace /PATH/TO/YOUR/DATA with the path of your training data, e.g., ../Data/OASIS, and make sure imgs and labels are properly loaded in the training script.

Step 2: Run python {training_script}, see "Training and testing scripts" for more details.

Publication

If you find this repository useful, please cite:

• Affine Medical Image Registration with Coarse-to-Fine Vision Transformer
  Tony C. W. Mok, Albert C. S. Chung
  CVPR2022. eprint arXiv:2203.15216

Acknowledgment

Some codes in this repository are modified from PVT and ViT. The MNI152 brain template is provided by the FLIRT (FMRIB's Linear Image Registration Tool).

Keywords

Keywords: Affine registration, Coarse-to-Fine Vision Transformer, 3D Vision Transformer