From 70edddf2527049f85e1cf27fa39108c7f17deeee Mon Sep 17 00:00:00 2001 From: Lin Tian Date: Wed, 31 Jul 2024 07:41:56 -0400 Subject: [PATCH 01/10] Add unittest for unigradicon-register. --- tests/test_command_arguments.py | 104 ++++++++++++++++++++++++++++++++ 1 file changed, 104 insertions(+) create mode 100644 tests/test_command_arguments.py diff --git a/tests/test_command_arguments.py b/tests/test_command_arguments.py new file mode 100644 index 0000000..92e7221 --- /dev/null +++ b/tests/test_command_arguments.py @@ -0,0 +1,104 @@ +import itk +import numpy as np +import unittest +import icon_registration.test_utils + +import subprocess +import os + + +class TestCommandInterface(unittest.TestCase): + def __init__(self, methodName: str = "runTest") -> None: + super().__init__(methodName) + icon_registration.test_utils.download_test_data() + self.test_data_dir = icon_registration.test_utils.TEST_DATA_DIR + self.test_temp_dir = f"{self.test_data_dir}/temp" + os.makedirs(self.test_temp_dir, exist_ok=True) + + def test_register_unigradicon_inference(self): + subprocess.run([ + "unigradicon-register", + "--fixed", f"{self.test_data_dir}/lung_test_data/copd1_highres_EXP_STD_COPD_img.nii.gz", + "--fixed_modality", "ct", + "--fixed_segmentation", f"{self.test_data_dir}/lung_test_data/copd1_highres_EXP_STD_COPD_label.nii.gz", + "--moving", f"{self.test_data_dir}/lung_test_data/copd1_highres_INSP_STD_COPD_img.nii.gz", + "--moving_modality", "ct", + "--moving_segmentation", f"{self.test_data_dir}/lung_test_data/copd1_highres_INSP_STD_COPD_label.nii.gz", + "--transform_out", f"{self.test_temp_dir}/transform.hdf5", + "--io_iterations", "None" + ]) + + # load transform + phi_AB = itk.transformread(f"{self.test_temp_dir}/transform.hdf5")[0] + + assert isinstance(phi_AB, itk.CompositeTransform) + + insp_points = icon_registration.test_utils.read_copd_pointset( + str( + icon_registration.test_utils.TEST_DATA_DIR + / "lung_test_data/copd1_300_iBH_xyz_r1.txt" + ) + ) + exp_points = icon_registration.test_utils.read_copd_pointset( + str( + icon_registration.test_utils.TEST_DATA_DIR + / "lung_test_data/copd1_300_eBH_xyz_r1.txt" + ) + ) + + dists = [] + for i in range(len(insp_points)): + px, py = ( + insp_points[i], + np.array(phi_AB.TransformPoint(tuple(exp_points[i]))), + ) + dists.append(np.sqrt(np.sum((px - py) ** 2))) + print(np.mean(dists)) + self.assertLess(np.mean(dists), 2.1) + + # remove temp file + os.remove(f"{self.test_temp_dir}/transform.hdf5") + + def test_register_unigradicon_io(self): + subprocess.run([ + "unigradicon-register", + "--fixed", f"{self.test_data_dir}/lung_test_data/copd1_highres_EXP_STD_COPD_img.nii.gz", + "--fixed_modality", "ct", + "--fixed_segmentation", f"{self.test_data_dir}/lung_test_data/copd1_highres_EXP_STD_COPD_label.nii.gz", + "--moving", f"{self.test_data_dir}/lung_test_data/copd1_highres_INSP_STD_COPD_img.nii.gz", + "--moving_modality", "ct", + "--moving_segmentation", f"{self.test_data_dir}/lung_test_data/copd1_highres_INSP_STD_COPD_label.nii.gz", + "--transform_out", f"{self.test_temp_dir}/transform.hdf5" + ]) + + # load transform + phi_AB = itk.transformread(f"{self.test_temp_dir}/transform.hdf5")[0] + + assert isinstance(phi_AB, itk.CompositeTransform) + + insp_points = icon_registration.test_utils.read_copd_pointset( + str( + icon_registration.test_utils.TEST_DATA_DIR + / "lung_test_data/copd1_300_iBH_xyz_r1.txt" + ) + ) + exp_points = icon_registration.test_utils.read_copd_pointset( + str( + icon_registration.test_utils.TEST_DATA_DIR + / "lung_test_data/copd1_300_eBH_xyz_r1.txt" + ) + ) + + dists = [] + for i in range(len(insp_points)): + px, py = ( + insp_points[i], + np.array(phi_AB.TransformPoint(tuple(exp_points[i]))), + ) + dists.append(np.sqrt(np.sum((px - py) ** 2))) + print(np.mean(dists)) + self.assertLess(np.mean(dists), 1.5) + + # remove temp file + os.remove(f"{self.test_temp_dir}/transform.hdf5") + From 6ed8200d5228066554f6e7da733a1852bc859780 Mon Sep 17 00:00:00 2001 From: Lin Tian Date: Wed, 31 Jul 2024 08:01:08 -0400 Subject: [PATCH 02/10] Enable the setting of the different similairty measure in the IO stage. --- src/unigradicon/__init__.py | 17 ++++++++++++++--- 1 file changed, 14 insertions(+), 3 deletions(-) diff --git a/src/unigradicon/__init__.py b/src/unigradicon/__init__.py index 4125e5e..a5b1113 100644 --- a/src/unigradicon/__init__.py +++ b/src/unigradicon/__init__.py @@ -159,9 +159,18 @@ def make_network(input_shape, include_last_step=False, lmbda=1.5, loss_fn=icon.L net.assign_identity_map(input_shape) return net +def make_sim(similarity): + if similarity == "lncc": + return icon.LNCC(sigma=5) + elif similarity == "lncc2": + return icon. SquaredLNCC(sigma=5) + elif similarity == "mind": + return icon.MINDSSC(radius=2, dilation=2) + else: + raise ValueError(f"Similarity measure {similarity} not recognized. Choose from [lncc, lncc2, mind].") -def get_unigradicon(): - net = make_network(input_shape, include_last_step=True) +def get_unigradicon(loss_fn=icon.LNCC(sigma=5)): + net = make_network(input_shape, include_last_step=True, loss_fn=loss_fn) from os.path import exists weights_location = "network_weights/unigradicon1.0/Step_2_final.trch" if not exists(weights_location): @@ -241,10 +250,12 @@ def main(): default=None, type=str, help="The path to save the warped image.") parser.add_argument("--io_iterations", required=False, default="50", help="The number of IO iterations. Default is 50. Set to 'None' to disable IO.") + parser.add_argument("--io_sim", required=False, + default="lncc", help="The similarity measure used in IO. Default is LNCC. Choose from [lncc, lncc2, mind].") args = parser.parse_args() - net = get_unigradicon() + net = get_unigradicon(make_sim(args.io_sim)) fixed = itk.imread(args.fixed) moving = itk.imread(args.moving) From 067d3ffb656dbeb010b61765a22c5911353db2dc Mon Sep 17 00:00:00 2001 From: Lin Tian Date: Wed, 31 Jul 2024 10:18:12 -0400 Subject: [PATCH 03/10] Enable user to load multigradicon. --- src/unigradicon/__init__.py | 32 ++++++++++++++++++++++++++++++-- tests/test_command_arguments.py | 14 ++++++++++---- 2 files changed, 40 insertions(+), 6 deletions(-) diff --git a/src/unigradicon/__init__.py b/src/unigradicon/__init__.py index a5b1113..d284ee5 100644 --- a/src/unigradicon/__init__.py +++ b/src/unigradicon/__init__.py @@ -169,12 +169,30 @@ def make_sim(similarity): else: raise ValueError(f"Similarity measure {similarity} not recognized. Choose from [lncc, lncc2, mind].") +def get_multigradicon(loss_fn=icon.LNCC(sigma=5)): + net = make_network(input_shape, include_last_step=True, loss_fn=loss_fn) + from os.path import exists + weights_location = "network_weights/multigradicon1.0/Step_2_final.trch" + if not exists(weights_location): + print("Downloading pretrained multigradicon model") + import urllib.request + import os + download_path = "https://github.com/uncbiag/uniGradICON/releases/download/multigradicon_weights/Step_2_final.trch" + os.makedirs("network_weights/multigradicon1.0/", exist_ok=True) + urllib.request.urlretrieve(download_path, weights_location) + print(f"Loading weights from {weights_location}") + trained_weights = torch.load(weights_location, map_location=torch.device("cpu")) + net.regis_net.load_state_dict(trained_weights) + net.to(config.device) + net.eval() + return net + def get_unigradicon(loss_fn=icon.LNCC(sigma=5)): net = make_network(input_shape, include_last_step=True, loss_fn=loss_fn) from os.path import exists weights_location = "network_weights/unigradicon1.0/Step_2_final.trch" if not exists(weights_location): - print("Downloading pretrained model") + print("Downloading pretrained unigradicon model") import urllib.request import os download_path = "https://github.com/uncbiag/uniGradICON/releases/download/unigradicon_weights/Step_2_final.trch" @@ -186,6 +204,14 @@ def get_unigradicon(loss_fn=icon.LNCC(sigma=5)): net.eval() return net +def get_model_from_model_zoo(model_name="unigradicon", loss_fn=icon.LNCC(sigma=5)): + if model_name == "unigradicon": + return get_unigradicon(loss_fn) + elif model_name == "multigradicon": + return get_multigradicon(loss_fn) + else: + raise ValueError(f"Model {model_name} not recognized. Choose from [unigradicon, multigradicon].") + def quantile(arr: torch.Tensor, q): arr = arr.flatten() l = len(arr) @@ -252,10 +278,12 @@ def main(): default="50", help="The number of IO iterations. Default is 50. Set to 'None' to disable IO.") parser.add_argument("--io_sim", required=False, default="lncc", help="The similarity measure used in IO. Default is LNCC. Choose from [lncc, lncc2, mind].") + parser.add_argument("--model", required=False, + default="unigradicon", help="The model to load. Default is unigradicon. Choose from [unigradicon, multigradicon].") args = parser.parse_args() - net = get_unigradicon(make_sim(args.io_sim)) + net = get_model_from_model_zoo(args.model, make_sim(args.io_sim)) fixed = itk.imread(args.fixed) moving = itk.imread(args.moving) diff --git a/tests/test_command_arguments.py b/tests/test_command_arguments.py index 92e7221..069839d 100644 --- a/tests/test_command_arguments.py +++ b/tests/test_command_arguments.py @@ -5,6 +5,7 @@ import subprocess import os +import torch class TestCommandInterface(unittest.TestCase): @@ -14,6 +15,7 @@ def __init__(self, methodName: str = "runTest") -> None: self.test_data_dir = icon_registration.test_utils.TEST_DATA_DIR self.test_temp_dir = f"{self.test_data_dir}/temp" os.makedirs(self.test_temp_dir, exist_ok=True) + self.device = torch.cuda.current_device() def test_register_unigradicon_inference(self): subprocess.run([ @@ -58,8 +60,8 @@ def test_register_unigradicon_inference(self): # remove temp file os.remove(f"{self.test_temp_dir}/transform.hdf5") - - def test_register_unigradicon_io(self): + + def test_register_multigradicon_inference(self): subprocess.run([ "unigradicon-register", "--fixed", f"{self.test_data_dir}/lung_test_data/copd1_highres_EXP_STD_COPD_img.nii.gz", @@ -68,7 +70,9 @@ def test_register_unigradicon_io(self): "--moving", f"{self.test_data_dir}/lung_test_data/copd1_highres_INSP_STD_COPD_img.nii.gz", "--moving_modality", "ct", "--moving_segmentation", f"{self.test_data_dir}/lung_test_data/copd1_highres_INSP_STD_COPD_label.nii.gz", - "--transform_out", f"{self.test_temp_dir}/transform.hdf5" + "--transform_out", f"{self.test_temp_dir}/transform.hdf5", + "--io_iterations", "None", + "--model", "multigradicon" ]) # load transform @@ -97,8 +101,10 @@ def test_register_unigradicon_io(self): ) dists.append(np.sqrt(np.sum((px - py) ** 2))) print(np.mean(dists)) - self.assertLess(np.mean(dists), 1.5) + self.assertLess(np.mean(dists), 3.8) # remove temp file os.remove(f"{self.test_temp_dir}/transform.hdf5") + + From f87f9db955a658d9c84399f556d17c6cda5740be Mon Sep 17 00:00:00 2001 From: Lin Tian Date: Wed, 31 Jul 2024 12:15:29 -0400 Subject: [PATCH 04/10] Add unit test workflow. --- .github/workflows/gpu-test-action.yml | 31 +++++++++++++++++++++++++++ 1 file changed, 31 insertions(+) create mode 100644 .github/workflows/gpu-test-action.yml diff --git a/.github/workflows/gpu-test-action.yml b/.github/workflows/gpu-test-action.yml new file mode 100644 index 0000000..5ab21de --- /dev/null +++ b/.github/workflows/gpu-test-action.yml @@ -0,0 +1,31 @@ +name: gpu-tests + +on: + pull_request: + push: + branches: [dev, main] + +jobs: + test-linux: + runs-on: [self-hosted, linux] + strategy: + max-parallel: 5 + + steps: + - uses: actions/checkout@v2 + - name: Set up Python 3.7 + uses: actions/setup-python@v2 + with: + python-version: 3.7 + - name: Install dependencies + run: | + pip install -r requirements.txt + + pip install -e . + + - name: fast test with unittest + run: | + python -m unittest -k CPU + - name: GPU test with unittest + run: | + python -m unittest discover \ No newline at end of file From c33a7970e469e1104e348108f34938eef9e6ac30 Mon Sep 17 00:00:00 2001 From: Lin Tian Date: Thu, 1 Aug 2024 12:41:39 -0400 Subject: [PATCH 05/10] 1. Add requirements.txt 2. Increase the version of icon_registration package to 1.1.5 --- requirements.txt | 1 + setup.cfg | 2 +- 2 files changed, 2 insertions(+), 1 deletion(-) create mode 100644 requirements.txt diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..994a988 --- /dev/null +++ b/requirements.txt @@ -0,0 +1 @@ +icon_registration>=1.1.5 \ No newline at end of file diff --git a/setup.cfg b/setup.cfg index a0575b5..dda77ff 100644 --- a/setup.cfg +++ b/setup.cfg @@ -21,7 +21,7 @@ packages = find: python_requires = >=3.7 install_requires = - icon_registration>=1.1.4 + icon_registration>=1.1.5 [options.packages.find] where = src From 89785e12a8b10ebe526b862c54f2feb9208530d7 Mon Sep 17 00:00:00 2001 From: HastingsGreer Date: Thu, 1 Aug 2024 13:50:38 -0400 Subject: [PATCH 06/10] Create test_requirements_sync.py --- tests/test_requirements_sync.py | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) create mode 100644 tests/test_requirements_sync.py diff --git a/tests/test_requirements_sync.py b/tests/test_requirements_sync.py new file mode 100644 index 0000000..f9ddb6b --- /dev/null +++ b/tests/test_requirements_sync.py @@ -0,0 +1,19 @@ +import unittest + + +class TestImports(unittest.TestCase): + + def test_requirements_match_cfg(self): + from inspect import getsourcefile + import os.path as path, sys + import configparser + + current_dir = path.dirname(path.abspath(getsourcefile(lambda: 0))) + parent_dir = current_dir[: current_dir.rfind(path.sep)] + + with open(parent_dir + "/requirements.txt") as f: + requirements_txt = "\n" + f.read() + requirements_cfg = configparser.ConfigParser() + requirements_cfg.read(parent_dir + "/setup.cfg") + requirements_cfg = requirements_cfg["options"]["install_requires"] + "\n" + self.assertEqual(requirements_txt, requirements_cfg) From 36876ea8dbca4f6a88a112f5165d5aac7d24efde Mon Sep 17 00:00:00 2001 From: HastingsGreer Date: Thu, 1 Aug 2024 13:57:26 -0400 Subject: [PATCH 07/10] Update test_requirements_sync.py --- tests/test_requirements_sync.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/test_requirements_sync.py b/tests/test_requirements_sync.py index f9ddb6b..d98922d 100644 --- a/tests/test_requirements_sync.py +++ b/tests/test_requirements_sync.py @@ -12,8 +12,8 @@ def test_requirements_match_cfg(self): parent_dir = current_dir[: current_dir.rfind(path.sep)] with open(parent_dir + "/requirements.txt") as f: - requirements_txt = "\n" + f.read() + requirements_txt = f.read() requirements_cfg = configparser.ConfigParser() requirements_cfg.read(parent_dir + "/setup.cfg") - requirements_cfg = requirements_cfg["options"]["install_requires"] + "\n" + requirements_cfg = requirements_cfg["options"]["install_requires"] self.assertEqual(requirements_txt, requirements_cfg) From 1b6b10616b65fd7d5b2a7442a825829e04aa0069 Mon Sep 17 00:00:00 2001 From: HastingsGreer Date: Thu, 1 Aug 2024 14:01:54 -0400 Subject: [PATCH 08/10] Update test_requirements_sync.py --- tests/test_requirements_sync.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_requirements_sync.py b/tests/test_requirements_sync.py index d98922d..f56b77a 100644 --- a/tests/test_requirements_sync.py +++ b/tests/test_requirements_sync.py @@ -12,7 +12,7 @@ def test_requirements_match_cfg(self): parent_dir = current_dir[: current_dir.rfind(path.sep)] with open(parent_dir + "/requirements.txt") as f: - requirements_txt = f.read() + requirements_txt = "\n" + f.read() requirements_cfg = configparser.ConfigParser() requirements_cfg.read(parent_dir + "/setup.cfg") requirements_cfg = requirements_cfg["options"]["install_requires"] From 63047714c56dba46760efedba54ee18944f1b0b6 Mon Sep 17 00:00:00 2001 From: Lin Tian Date: Mon, 5 Aug 2024 04:51:02 -0400 Subject: [PATCH 09/10] Increase version number. --- setup.cfg | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/setup.cfg b/setup.cfg index dda77ff..f50b557 100644 --- a/setup.cfg +++ b/setup.cfg @@ -1,8 +1,8 @@ [metadata] name = unigradicon -version = 1.0.2 +version = 1.0.3 author = Lin Tian -author_email = +author_email = lintian@cs.unc.edu description = a foundation model for medical image registration long_description = file: README.md long_description_content_type = text/markdown From 4960f0c65b4d40515da94c4d07d1515aaf6bd905 Mon Sep 17 00:00:00 2001 From: Lin Tian Date: Mon, 5 Aug 2024 05:17:25 -0400 Subject: [PATCH 10/10] Update README to include multiGradICON reference and usages. --- README.md | 37 ++++++++++++++++++++++++++++--------- 1 file changed, 28 insertions(+), 9 deletions(-) diff --git a/README.md b/README.md index eb9ce29..b25fe76 100644 --- a/README.md +++ b/README.md @@ -11,13 +11,18 @@ The result is a deep-learning-based registration model that works well across da Please (currently) cite as: ``` -@misc{tian2024unigradicon, - title={uniGradICON: A Foundation Model for Medical Image Registration}, - author={Lin Tian and Hastings Greer and Roland Kwitt and Francois-Xavier Vialard and Raul San Jose Estepar and Sylvain Bouix and Richard Rushmore and Marc Niethammer}, - year={2024}, - eprint={2403.05780}, - archivePrefix={arXiv}, - primaryClass={cs.CV} +@article{tian2024unigradicon, + title={uniGradICON: A Foundation Model for Medical Image Registration}, + author={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}, + journal={arXiv preprint arXiv:2403.05780}, + year={2024} +} + +@article{demir2024multigradicon, + title={multiGradICON: A Foundation Model for Multimodal Medical Image Registration}, + author={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}, + journal={arXiv preprint arXiv:2408.00221}, + year={2024} } ``` @@ -204,12 +209,25 @@ unigradicon-register --fixed=RegLib_C01_2.nrrd --fixed_modality=mri --moving=Reg ``` -To register without instance optimization +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 warp +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 load specific model weight in the 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 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 ``` @@ -218,6 +236,7 @@ To warp a label map ``` 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 ``` + We also provide a [colab](https://colab.research.google.com/drive/1JuFL113WN3FHCoXG-4fiBTWIyYpwGyGy?usp=sharing) demo.