From f7a6819fc5553b1ec16a5d49a2b06edac343830e Mon Sep 17 00:00:00 2001
From: IgorTatarnikov <igor.tatarnikov@gmail.com>
Date: Tue, 4 Jun 2024 10:52:17 +0100
Subject: [PATCH] WIP calculate inverse transform

---
 brainglobe_registration/elastix/register.py   | 54 +++++++++++--------
 .../registration_widget.py                    | 14 ++++-
 2 files changed, 44 insertions(+), 24 deletions(-)

diff --git a/brainglobe_registration/elastix/register.py b/brainglobe_registration/elastix/register.py
index 607a059..481eb90 100644
--- a/brainglobe_registration/elastix/register.py
+++ b/brainglobe_registration/elastix/register.py
@@ -1,4 +1,4 @@
-from typing import List
+from typing import Dict, List, Tuple
 
 import itk
 import numpy as np
@@ -28,8 +28,10 @@ def run_registration(
     atlas_image,
     moving_image,
     annotation_image,
-    parameter_lists: List[tuple[str, dict]],
-) -> tuple[np.ndarray, itk.ParameterObject, np.ndarray]:
+    parameter_lists: List[Tuple[str, Dict]],
+) -> Tuple[
+    np.ndarray, itk.ParameterObject, np.ndarray, np.ndarray, np.ndarray
+]:
     """
     Run the registration process on the given images.
 
@@ -41,7 +43,7 @@ def run_registration(
         The moving image.
     annotation_image : np.ndarray
         The annotation image.
-    parameter_lists : List[tuple[str, dict]], optional
+    parameter_lists : List[Tuple[str, Dict]], optional
         The list of parameter lists, by default None
 
     Returns
@@ -90,41 +92,47 @@ def run_registration(
 
     # Invert the transformation
     # Import Default Parameter Map and adjust parameters
-    parameter_object_inverse = itk.ParameterObject.New()
-    parameter_map_rigid = parameter_object_inverse.GetDefaultParameterMap(
-        "rigid"
-    )
-    parameter_map_bspline = parameter_object_inverse.GetDefaultParameterMap(
-        "bspline"
-    )
-    parameter_map_bspline["HowToCombineTransforms"] = ["Compose"]
-    parameter_map_rigid["HowToCombineTransforms"] = ["Compose"]
-    parameter_object_inverse.AddParameterMap(parameter_map_rigid)
-    parameter_object_inverse.AddParameterMap(parameter_map_bspline)
+    parameter_object_inverse = setup_parameter_object(parameter_lists[-1::-1])
+    for i in range(len(parameter_lists)):
+        parameter_object_inverse.SetParameter(
+            i, "HowToCombineTransforms", ["Compose"]
+        )
 
     (
         inverse_image,
         inverse_transform_parameters,
     ) = itk.elastix_registration_method(
-        atlas_image_elastix,
-        atlas_image_elastix,
-        parameter_object=parameter_object,
-        initial_transform_parameter_file_name="output/TransformParameters.2.txt",
+        moving_image_elastix,
+        moving_image_elastix,
+        parameter_object=parameter_object_inverse,
+        initial_transform_parameter_file_name=f"output/TransformParameters.{len(parameter_lists)-1}.txt",
     )
 
     # Adjust inverse transform parameters object
     inverse_transform_parameters.SetParameter(
-        0, "InitialTransformParametersFileName", "NoInitialTransform"
+        0, "InitialTransformParameterFileName", "NoInitialTransform"
+    )
+
+    file_names = [
+        f"InverseTransformParameters.{i}.txt"
+        for i in range(len(parameter_lists))
+    ]
+
+    itk.ParameterObject.WriteParameterFiles(
+        inverse_transform_parameters, file_names
     )
 
-    # inverse_transform_parameters.WriteParameterFiles(
-    # "output/TransformParameters_inverse"
-    # )
+    inverse_moving = itk.transformix_filter(
+        moving_image,
+        inverse_transform_parameters,
+    )
 
     return (
         np.asarray(result_image),
         result_transform_parameters,
         np.asarray(annotation_image_transformix),
+        np.asarray(inverse_image),
+        np.asarray(inverse_moving),
     )
 
 
diff --git a/brainglobe_registration/registration_widget.py b/brainglobe_registration/registration_widget.py
index 67ab52f..de72798 100644
--- a/brainglobe_registration/registration_widget.py
+++ b/brainglobe_registration/registration_widget.py
@@ -205,12 +205,24 @@ def _on_adjust_moving_image_reset_button_click(self):
     def _on_run_button_click(self):
         current_atlas_slice = self._viewer.dims.current_step[0]
 
-        result, parameters, registered_annotation_image = run_registration(
+        (
+            result,
+            parameters,
+            registered_annotation_image,
+            inverse_image,
+            inverse_test,
+        ) = run_registration(
             self._atlas.reference[current_atlas_slice, :, :],
             self._moving_image.data,
             self._atlas.annotation[current_atlas_slice, :, :],
             self.transform_selections,
         )
+        self._viewer.add_image(
+            inverse_image, name="Inverse Image", visible=False
+        )
+        self._viewer.add_image(
+            inverse_test, name="Inverse Test", visible=False
+        )
 
         boundaries = find_boundaries(
             registered_annotation_image, mode="inner"