ICON-console-scripts

bin/icon_register_pair

@@ -0,0 +1,71 @@
+#!/usr/bin/env python
+
+import argparse
+
+import icon_registration as icon
+import icon_registration.itk_wrapper
+import itk
+import numpy as np
+import torch
+import importlib
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--model", required=True)
+parser.add_argument("--weights_path")
+parser.add_argument("--fixed_image", required=True)
+parser.add_argument("--moving_image", required=True)
+parser.add_argument("--transform_out")
+parser.add_argument("--displacement_image_out")
+parser.add_argument("--warped_moving_out")
+parser.add_argument("--finetune", action='store_true')
+
+args = parser.parse_args()
+
+model_package_name = ".".join(args.model.split(".")[:-1])
+model_name = args.model.split(".")[-1]
+
+model_package = importlib.import_module(model_package_name)
+net = getattr(model_package, model_name)()
+
+if args.weights_path:
+ weights = torch.load(args.weights_path)
+ net.regis_net.load_state_dict(weights)
+
+fixed_image = itk.imread(args.fixed_image)
+moving_image = itk.imread(args.moving_image)
+
+# We want images normalized to 0.0, 1.0
+for image in fixed_image, moving_image:
+ assert(np.abs(np.max(np.array(image)) - 1.0) < 0.2)
+ assert(np.abs(np.min(np.array(image))) < 0.2)
+
+phi, _ = icon.itk_wrapper.register_pair(
+ net, moving_image, fixed_image, finetune_steps=90 if args.finetune else None
+)
+
+if args.transform_out:
+ itk.transformwrite([phi], args.transform_out)
+
+if args.displacement_image_out:
+ filter = itk.TransformToDisplacementFieldFilter[
+ itk.itkImagePython.itkImageVF33, itk.D
+ ].New()
+ decorator = itk.DataObjectDecorator[itk.Transform[itk.D, 3, 3]].New()
+ decorator.Set(phi)
+ filter.SetInput(decorator)
+ filter.SetReferenceImage(fixed_image)
+ filter.SetUseReferenceImage(True)
+ filter.Update()
+ itk.imwrite(filter.GetOutput(), args.displacement_image_out)
+
+if args.warped_moving_out:
+ warped = itk.resample_image_filter(
+ moving_image,
+ use_reference_image=True,
+ reference_image=fixed_image,
+ transform=phi,
+ interpolator=itk.LinearInterpolateImageFunction.New(moving_image),
+ )
+ itk.imwrite(warped, args.warped_moving_out)
+
+

setup.cfg

@@ -20,6 +20,9 @@ package_dir =
 packages = find:
 python_requires = >=3.4
 
+scripts = 
+ bin/icon_register_pair
+
 install_requires =
  torch
  torchvision

test/test_console_scripts.py

@@ -0,0 +1,23 @@
+import unittest
+import subprocess
+from icon_registration import test_utils
+
+class TestConsoleScripts(unittest.TestCase):
+ def test_register_pair_script(self):
+ import footsteps
+ test_utils.download_test_data()
+ subprocess.run(
+ [
+ "icon_register_pair",
+ "--fixed_image", 
+ test_utils.TEST_DATA_DIR / "brain_test_data" / "2_T1w_acpc_dc_restore_brain.nii.gz",
+ "--moving_image", 
+ test_utils.TEST_DATA_DIR / "brain_test_data" / "8_T1w_acpc_dc_restore_brain.nii.gz",
+ "--model",
+ "icon_registration.pretrained_models.brain_registration_model"
+ "--warped_image_out", 
+ footsteps.output_dir + "warped.nii.gz"
+ ]
+ )
+ import itk
+ itk.imread(footsteps.output_dir + "warped.nii.gz")