From 3f4a0479913722981759b6643138682bb1b35099 Mon Sep 17 00:00:00 2001 From: Lin Tian Date: Tue, 7 May 2024 00:55:49 -0400 Subject: [PATCH] Add reorient functionality to main() function. --- src/unigradicon/__init__.py | 19 +++++++++++++++++-- 1 file changed, 17 insertions(+), 2 deletions(-) diff --git a/src/unigradicon/__init__.py b/src/unigradicon/__init__.py index 4125e5e..2ed35a1 100644 --- a/src/unigradicon/__init__.py +++ b/src/unigradicon/__init__.py @@ -217,6 +217,14 @@ def preprocess(image, modality="ct", segmentation=None): image = apply_mask(image, segmentation) return image +def reorient(moving, fixed): + from itk.ITKCommonBasePython import itkSpatialOrientationAdapter + + return itk.orient_image_filter( + moving, + desired_coordinate_orientation=itkSpatialOrientationAdapter().FromDirectionCosines(fixed.GetDirection()), + use_image_direction=True) + def main(): import itk import argparse @@ -241,6 +249,7 @@ 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("--reorient", action="store_true", help="Reorient the source image to the orientation of the target image.") args = parser.parse_args() @@ -264,10 +273,16 @@ def main(): else: io_iterations = int(args.io_iterations) + moving_processed = preprocess(moving, args.moving_modality, moving_segmentation) + fixed_processed = preprocess(fixed, args.fixed_modality, fixed_segmentation) + + if args.reorient: + moving_processed = reorient(moving_processed, fixed_processed) + phi_AB, phi_BA = icon_registration.itk_wrapper.register_pair( net, - preprocess(moving, args.moving_modality, moving_segmentation), - preprocess(fixed, args.fixed_modality, fixed_segmentation), + moving_processed, + fixed_processed, finetune_steps=io_iterations) itk.transformwrite([phi_AB], args.transform_out)