Merge pull request #119 from sarthakpati/greedy

Added `greedy` and `elastix` registrators

Author: MarcelRosier

brainles_preprocessing/registration/__init__.py

@@ -17,3 +17,7 @@
 
 
 from .niftyreg.niftyreg import NiftyRegRegistrator
+
+from .elastix.elastix import ElastixRegistrator
+
+from .greedy.greedy import GreedyRegistrator

brainles_preprocessing/registration/elastix/__init__.py

@@ -0,0 +1,119 @@
+# TODO add typing and docs
+from typing import Optional
+import os
+
+import itk
+
+from brainles_preprocessing.registration.registrator import Registrator
+from brainles_preprocessing.utils import check_and_add_suffix
+
+
+class ElastixRegistrator(Registrator):
+    def __init__(
+        self,
+    ):
+        pass
+
+    def register(
+        self,
+        fixed_image_path: str,
+        moving_image_path: str,
+        transformed_image_path: str,
+        matrix_path: str,
+        log_file_path: Optional[str] = None,
+        parameter_object: Optional[itk.ParameterObject] = None,
+    ) -> None:
+        """
+        Register images using elastix.
+
+        Args:
+            fixed_image_path (str): Path to the fixed image.
+            moving_image_path (str): Path to the moving image.
+            transformed_image_path (str): Path to the transformed image (output).
+            matrix_path (str): Path to the transformation matrix (output). This gets overwritten if it already exists.
+            log_file_path (Optional[str]): Path to the log file.
+            parameter_object (Optional[itk.ParameterObject]): The parameter object for elastix registration.
+        """
+        # initialize parameter object
+        if parameter_object is None:
+            parameter_object = self.__initialize_parameter_object()
+        # add .txt suffix to the matrix path if it doesn't have any extension
+        matrix_path = check_and_add_suffix(matrix_path, ".txt")
+
+        # read images as itk images
+        fixed_image = itk.imread(fixed_image_path)
+        moving_image = itk.imread(moving_image_path)
+
+        if log_file_path is not None:
+            # split log_file_path
+            log_path, log_file = os.path.split(log_file_path)
+            result_image, result_transform_params = itk.elastix_registration_method(
+                fixed_image,
+                moving_image,
+                parameter_object=parameter_object,
+                log_to_file=True,
+                log_file_name=log_file,
+                output_directory=log_path,
+            )
+        else:
+            result_image, result_transform_params = itk.elastix_registration_method(
+                fixed_image,
+                moving_image,
+                parameter_object=parameter_object,
+                log_to_console=True,
+            )
+
+        itk.imwrite(result_image, transformed_image_path)
+
+        if not os.path.exists(matrix_path):
+            result_transform_params.WriteParameterFile(
+                result_transform_params.GetParameterMap(0),
+                matrix_path,
+            )
+
+    def transform(
+        self,
+        fixed_image_path: str,
+        moving_image_path: str,
+        transformed_image_path: str,
+        matrix_path: str,
+        log_file_path: Optional[str] = None,
+    ) -> None:
+        """
+        Apply a transformation using elastix.
+
+        Args:
+            fixed_image_path (str): Path to the fixed image.
+            moving_image_path (str): Path to the moving image.
+            transformed_image_path (str): Path to the transformed image (output).
+            matrix_path (str): Path to the transformation matrix (output). This gets overwritten if it already exists.
+            log_file_path (Optional[str]): Path to the log file.
+        """
+        parameter_object = self.__initialize_parameter_object()
+
+        # check if the matrix file exists
+        if os.path.exists(matrix_path):
+            parameter_object.SetParameter(
+                0, "InitialTransformParametersFileName", matrix_path
+            )
+
+        self.register(
+            fixed_image_path,
+            moving_image_path,
+            transformed_image_path,
+            matrix_path,
+            log_file_path,
+            parameter_object,
+        )
+
+    def __initialize_parameter_object(self) -> itk.ParameterObject:
+        """
+        Initialize the parameter object for elastix registration.
+
+        Returns:
+            itk.ParameterObject: The parameter object for registration.
+        """
+        parameter_object = itk.ParameterObject.New()
+        default_rigid_parameter_map = parameter_object.GetDefaultParameterMap("rigid")
+        parameter_object.AddParameterMap(default_rigid_parameter_map)
+        return parameter_object

brainles_preprocessing/registration/elastix/elastix.py

@@ -0,0 +1,96 @@
+# TODO add typing and docs
+from typing import Optional
+import contextlib
+import os
+
+from picsl_greedy import Greedy3D
+
+from brainles_preprocessing.registration.registrator import Registrator
+from brainles_preprocessing.utils import check_and_add_suffix
+
+
+class GreedyRegistrator(Registrator):
+    def __init__(
+        self,
+    ):
+        pass
+
+    def register(
+        self,
+        fixed_image_path: str,
+        moving_image_path: str,
+        transformed_image_path: str,
+        matrix_path: str,
+        log_file_path: Optional[str] = None,
+    ) -> None:
+        """
+        Register images using greedy. Ref: https://pypi.org/project/picsl-greedy/ and https://greedy.readthedocs.io/en/latest/reference.html#greedy-usage
+
+        Args:
+            fixed_image_path (str): Path to the fixed image.
+            moving_image_path (str): Path to the moving image.
+            transformed_image_path (str): Path to the transformed image (output).
+            matrix_path (str): Path to the transformation matrix (output). This gets overwritten if it already exists.
+            log_file_path (Optional[str]): Path to the log file, which is not used.
+        """
+        # add .txt suffix to the matrix path if it doesn't have any extension
+        matrix_path = check_and_add_suffix(matrix_path, ".mat")
+
+        registor = Greedy3D()
+        # these parameters are taken from the OG BraTS Pipeline [https://github.com/CBICA/CaPTk/blob/master/src/applications/BraTSPipeline.cxx]
+        command_to_run = f"-i {fixed_image_path} {moving_image_path} -o {matrix_path} -a -dof 6 -m NMI -n 100x50x5 -ia-image-centers"
+
+        if log_file_path is not None:
+            with open(log_file_path, "a+") as f:
+                with contextlib.redirect_stdout(f):
+                    registor.execute(command_to_run)
+        else:
+            registor.execute(command_to_run)
+
+        self.transform(
+            fixed_image_path, moving_image_path, transformed_image_path, matrix_path
+        )
+
+    def transform(
+        self,
+        fixed_image_path: str,
+        moving_image_path: str,
+        transformed_image_path: str,
+        matrix_path: str,
+        interpolator: Optional[str] = "LINEAR",
+        log_file_path: Optional[str] = None,
+    ) -> None:
+        """
+        Apply a transformation using greedy.
+
+        Args:
+            fixed_image_path (str): Path to the fixed image.
+            moving_image_path (str): Path to the moving image.
+            transformed_image_path (str): Path to the transformed image (output).
+            matrix_path (str): Path to the transformation matrix (output). This gets overwritten if it already exists.
+            interpolator (Optional[str]): The interpolator to use; one of NN, LINEAR or LABEL.
+            log_file_path (Optional[str]): Path to the log file, which is not used.
+        """
+        registor = Greedy3D()
+        interpolator_upper = interpolator.upper()
+        if "LABEL" in interpolator_upper:
+            interpolator_upper += " 0.3vox"
+
+        matrix_path = check_and_add_suffix(matrix_path, ".mat")
+
+        if not os.path.exists(matrix_path):
+            self.register(
+                fixed_image_path,
+                moving_image_path,
+                transformed_image_path,
+                matrix_path,
+                log_file_path,
+            )
+
+        command_to_run = f"-rf {fixed_image_path} -rm {moving_image_path} {transformed_image_path} -r {matrix_path} -ri {interpolator_upper}"
+        if log_file_path is not None:
+            with open(log_file_path, "a+") as f:
+                with contextlib.redirect_stdout(f):
+                    registor.execute(command_to_run)
+        else:
+            registor.execute(command_to_run)

brainles_preprocessing/registration/greedy/__init__.py

@@ -0,0 +1 @@
+from .generic import check_and_add_suffix

brainles_preprocessing/registration/greedy/greedy.py

@@ -0,0 +1,15 @@
+def check_and_add_suffix(filename: str, suffix: str) -> str:
+    """
+    Adds a suffix to the filename if it doesn't already have it.
+
+    Parameters:
+        filename (str): The filename to check and potentially modify.
+        suffix (str): The suffix to add to the filename.
+
+    Returns:
+        str: The filename with the suffix added if needed.
+    """
+    filename_copy = filename
+    if not filename_copy.endswith(suffix):
+        filename_copy += suffix
+    return filename_copy

brainles_preprocessing/utils/__init__.py

@@ -70,11 +70,15 @@ rich = "^13.6.0"
 
 # optional registration backends
 ereg = { version = "^0.0.10", optional = true }
+itk-elastix = { version = "^0.20.0", optional = true }
+picsl_greedy = { version = "^0.0.6", optional = true }
 
 
 [tool.poetry.extras]
-all = ["ereg"]
+all = ["ereg", "itk-elastix", "picsl_greedy"]
 ereg = ["ereg"]
+itk-elastix = ["itk-elastix"]
+picsl_greedy = ["picsl_greedy"]
 
 
 [tool.poetry.group.dev.dependencies]

brainles_preprocessing/utils/generic.py

@@ -3,6 +3,8 @@
 from brainles_preprocessing.registration.ANTs.ANTs import ANTsRegistrator
 from brainles_preprocessing.registration.eReg.eReg import eRegRegistrator
 from brainles_preprocessing.registration.niftyreg.niftyreg import NiftyRegRegistrator
+from brainles_preprocessing.registration.elastix.elastix import ElastixRegistrator
+from brainles_preprocessing.registration.greedy.greedy import GreedyRegistrator
 
 import unittest
 
@@ -29,3 +31,19 @@ def get_registrator(self):
 
     def get_method_and_extension(self):
         return "ereg", "mat"
+
+
+class TestElastixRegistrator(RegistratorBase, unittest.TestCase):
+    def get_registrator(self):
+        return ElastixRegistrator()
+
+    def get_method_and_extension(self):
+        return "elastix", "txt"
+
+
+class TestGreedyRegistrator(RegistratorBase, unittest.TestCase):
+    def get_registrator(self):
+        return GreedyRegistrator()
+
+    def get_method_and_extension(self):
+        return "greedy", "mat"