uniGradICON: A Foundation Model for Medical Image Registration

News

• 🔥 Finetuning code for uniGradICON and multiGradICON is released (10/2026)
• 🔥 uniGradICON Slicer Extension is available in 3D Slicer 5.7.0 Extensions Manager (12/2024)
• 🏆 multiGradICON wins the best oral presentation at 2024 MICCAI Workshop for Biomedical Image Registration (WBIR) (10/2024)
• 🔥 uniGradICON has been used as a baseline in LUMIR Brain MRI Registration Challenge (6/2024)

Introduction

uniGradICON is based on GradICON but trained on several different datasets (see details below). The result is a deep-learning-based registration model that works well across datasets. More results can be found here. To get you started quickly there is also an easy to use Slicer extension.

This is the repository for the following papers

uniGradICON: A Foundation Model for Medical Image Registration
Tian, Lin and Greer, Hastings and Kwitt, Roland and Vialard, Francois-Xavier and Estepar, Raul San Jose and Bouix, Sylvain and Rushmore, Richard and Niethammer, Marc
MICCAI 2024 https://arxiv.org/abs/2403.05780

multiGradICON: A Foundation Model for Multimodal Medical Image Registration
Demir, Basar and Tian, Lin and Greer, Thomas Hastings and Kwitt, Roland and Vialard, Francois-Xavier and Estepar, Raul San Jose and Bouix, Sylvain and Rushmore, Richard Jarrett and Ebrahim, Ebrahim and Niethammer, Marc
MICCAI Workshop on Biomedical Image Registration (WBIR) 2024 https://arxiv.org/abs/2408.00221

Easy to use and install

The pre-trained uniGradICON and multiGradICON can be used via CLI, colab notebook, and Slicer Extension. The model weights will be downloaded automatically. You can also find the model weights here.

👉 Inference via CLI

Installation

python3 -m venv unigradicon_virtualenv
source unigradicon_virtualenv/bin/activate

pip install unigradicon

To register one pair of images

wget https://www.hgreer.com/assets/slicer_mirror/RegLib_C01_1.nrrd
wget https://www.hgreer.com/assets/slicer_mirror/RegLib_C01_2.nrrd

unigradicon-register --fixed=RegLib_C01_2.nrrd --fixed_modality=mri --moving=RegLib_C01_1.nrrd --moving_modality=mri --transform_out=trans.hdf5 --warped_moving_out=warped_C01_1.nrrd

To register without instance optimization (IO)

unigradicon-register --fixed=RegLib_C01_2.nrrd --fixed_modality=mri --moving=RegLib_C01_1.nrrd --moving_modality=mri --transform_out=trans.hdf5 --warped_moving_out=warped_C01_1.nrrd --io_iterations None

To use a different similarity measure in the IO. We currently support three similarity measures

• LNCC: lncc
• Squared LNCC: lncc2
• MIND SSC: mind

unigradicon-register --fixed=RegLib_C01_2.nrrd --fixed_modality=mri --moving=RegLib_C01_1.nrrd --moving_modality=mri --transform_out=trans.hdf5 --warped_moving_out=warped_C01_1.nrrd --io_iterations 50 --io_sim lncc2

To change the learning rate of the IO (default 0.00002)

unigradicon-register --fixed=RegLib_C01_2.nrrd --fixed_modality=mri --moving=RegLib_C01_1.nrrd --moving_modality=mri --transform_out=trans.hdf5 --warped_moving_out=warped_C01_1.nrrd --io_iterations 50 --io_lr 0.00002

To load specific model weights during inference. We currently support uniGradICON and multiGradICON.

unigradicon-register --fixed=RegLib_C01_2.nrrd --fixed_modality=mri --moving=RegLib_C01_1.nrrd --moving_modality=mri --transform_out=trans.hdf5 --warped_moving_out=warped_C01_1.nrrd --model multigradicon

To mask out the background using the provided masks before registration (masks for both moving and fixed images are necessary for accurate registration):

unigradicon-register --fixed=RegLib_C01_2.nrrd --fixed_modality=mri --fixed_mask=[fixed_image_mask_file_name] --moving=RegLib_C01_1.nrrd --moving_modality=mri --moving_mask=[moving_image_mask_file_name] --transform_out=trans.hdf5 --warped_moving_out=warped_C01_1.nrrd --io_iterations None --input_masking

To apply loss function masking using the provided masks in the IO:

unigradicon-register --fixed=RegLib_C01_2.nrrd --fixed_modality=mri --fixed_mask=[fixed_image_mask_file_name] --moving=RegLib_C01_1.nrrd --moving_modality=mri --moving_mask=[moving_image_mask_file_name] --transform_out=trans.hdf5 --warped_moving_out=warped_C01_1.nrrd --io_iterations 50 --io_sim lncc2 --loss_function_masking

Masks must be on the same voxel grid as their image and saved in the same format, for example both as NIfTI.

As in 1.0.5, segmentations are still used as masks when no masks and no Dice loss are given. This is deprecated, so please use --fixed_mask and --moving_mask instead.

Intensity conservation loss ensures proper intensity adjustments for registration tasks requiring mass conservation, utilizing the change of variables rule from integration. To be effective, images must be in an intensity space where conservation holds. This loss function is specifically valid for CT modality, where -1000 HU represents air. To apply determinant-based intensity correction during registration in the IO:

unigradicon-register --fixed=[fixed_ct_file_name] --fixed_modality=ct --moving=[moving_ct_file_name] --moving_modality=ct --transform_out=trans.hdf5 --warped_moving_out=warped_ct.nii.gz --io_iterations 50 --io_sim lncc2 --intensity_conservation_loss

To apply Dice loss using the provided segmentations in the IO (both segmentations must share at least one non-zero label)

unigradicon-register --fixed=RegLib_C01_2.nrrd --fixed_modality=mri --fixed_segmentation=[fixed_image_segmentation_file_name] --moving=RegLib_C01_1.nrrd --moving_modality=mri --moving_segmentation=[moving_image_segmentation_file_name] --transform_out=trans.hdf5 --warped_moving_out=warped_C01_1.nrrd --io_iterations 50 --io_sim lncc2 --dice_loss_weight 0.1

To finetune on your own data, see the finetuning guide.

unigradicon-finetune --config /path/to/config.yaml

To use custom network weights, for example after finetuning

unigradicon-register --fixed=fixed.nii.gz --fixed_modality=mri --moving=moving.nii.gz --moving_modality=mri --transform_out=trans.hdf5 --warped_moving_out=warped.nii.gz --network_weights /path/to/network_weights_final.trch

If you customized preprocessing during finetuning, pass the same values at inference with --ct_window (for CT) or --quantile_range (for MRI).

unigradicon-register --fixed=fixed.nii.gz --fixed_modality=mri --moving=moving.nii.gz --moving_modality=mri --transform_out=trans.hdf5 --quantile_range 0.01 0.95 --network_weights /path/to/network_weights_final.trch

To warp an image

unigradicon-warp --fixed [fixed_image_file_name] --moving [moving_image_file_name]  --transform trans.hdf5 --warped_moving_out warped.nii.gz --linear

To warp a segmentation

unigradicon-warp --fixed [fixed_image_file_name] --moving [moving_image_segmentation_file_name]  --transform trans.hdf5 --warped_moving_out warped_seg.nii.gz --nearest_neighbor

👉 Inference via colab notebook

We provide a colab notebook where the users can directly access and visualize the output of uniGradICON network.

👉 Inference via Slicer Extension

A Slicer extension is available here. It is an official Slicer Extension and can be installed via the Slicer Extension Manager. This requires Slicer >=5.7.0. Please make sure to install the Slicer PyTorch extension first, since uniGradICON depends on it.

Training and testing data

uniGradICON has currently been trained and tested on the following datasets.

Training data:

	Dataset	Anatomical region	# of patients	# per patient	# of pairs	Type	Modality
1.	COPDGene	Lung	899	2	899	Intra-pat.	CT
2.	OAI	Knee	2532	1	3,205,512	Inter-pat.	MRI
3.	HCP	Brain	1076	1	578,888	Inter-pat.	MRI
4.	L2R-Abdomen	Abdomen	30	1	450	Inter-pat.	CT

Testing data:

	Dataset	Anatomical region	# of patients	# per patient	# of pairs	Type	Modality
5.	Dirlab-COPDGene	Lung	10	2	10	Intra-pat.	CT
6.	OAI-test	Knee	301	1	301	Inter-pat.	MRI
7.	HCP-test	Brain	32	1	100	Inter-pat.	MRI
8.	L2R-NLST-val	Lung	10	2	10	Intra-pat.	CT
9.	L2R-OASIS-val	Brain	20	1	19	Inter-pat.	MRI
10.	IXI-test	Brain	115	1	115	Atlas-pat.	MRI
11.	L2R-CBCT-val	Lung	3	3	6	Intra-pat.	CT/CBCT
12.	L2R-CTMR-val	Abdomen	3	2	3	Intra-pat.	CT/MRI
13.	L2R-CBCT-train	Lung	3	11	22	Intra-pat.	CT/CBCT

Finetuning

You can finetune uniGradICON or multiGradICON on your own data using the finetuning guide.

Get involved

Our goal is to continuously improve the uniGradICON model, e.g., by training on more datasets with additional diversity. Feel free to point us to datasets that should be included or let us know if you want to help with future developments.

Visualization

uniGradICON is set up to work with ITK images and transforms. So you can easily read and write images and display resulting transformations for example in 3D Slicer.

The result can be viewed in 3D Slicer:

📄 Citation

If you find this repository useful, please consider citing:

@inproceedings{tian2024unigradicon,
  title={unigradicon: A foundation model for medical image registration},
  author={Tian, Lin and Greer, Hastings and Kwitt, Roland and Vialard, Fran{\c{c}}ois-Xavier and San Jos{\'e} Est{\'e}par, Ra{\'u}l and Bouix, Sylvain and Rushmore, Richard and Niethammer, Marc},
  booktitle={International Conference on Medical Image Computing and Computer-Assisted Intervention},
  pages={749--760},
  year={2024},
  organization={Springer}
}

@inproceedings{demir2024multigradicon,
  title={MultiGradICON: A foundation model for multimodal medical image registration},
  author={Demir, Ba{\c{s}}ar and Tian, Lin and Greer, Hastings and Kwitt, Roland and Vialard, Fran{\c{c}}ois-Xavier and Est{\'e}par, Ra{\'u}l San Jos{\'e} and Bouix, Sylvain and Rushmore, Richard and Ebrahim, Ebrahim and Niethammer, Marc},
  booktitle={International Workshop on Biomedical Image Registration},
  pages={3--18},
  year={2024},
  organization={Springer}
}