Merge pull request #49 from BrainLesion/introduce_ereg

Introduce_ereg

brainles_preprocessing/registration/__init__.py

@@ -7,4 +7,12 @@
         "ANTS package not found. If you want to use it, please install it using 'pip install brainles_preprocessing[ants]'"
     )
 
+try:
+    from .eReg.eReg import eRegRegistrator
+except ImportError:
+    warnings.warn(
+        "eReg package not found. If you want to use it, please install it using 'pip install brainles_preprocessing[ereg]'"
+    )
+
+
 from .niftyreg.niftyreg import NiftyRegRegistrator

brainles_preprocessing/registration/eReg/eReg.py

@@ -0,0 +1,100 @@
+# TODO add typing and docs
+import os
+
+from ereg.registration import RegistrationClass
+
+from brainles_preprocessing.registration.registrator import Registrator
+
+
+class eRegRegistrator(Registrator):
+    def __init__(
+        self,
+        # TODO define default
+        configuration_file: str | None = None,
+    ):
+        """
+        # TODO
+        """
+        self.configuration_file = configuration_file
+
+    def register(
+        self,
+        fixed_image_path: str,
+        moving_image_path: str,
+        transformed_image_path: str,
+        matrix_path: str,
+        log_file_path: str = None,
+    ) -> None:
+        """
+        Register images using eReg.
+
+        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).
+            log_file_path (str): Path to the log file.
+        """
+        # TODO do we need to handle kwargs?
+        registrator = RegistrationClass(
+            configuration_file=self.configuration_file,
+        )
+
+        matrix_path = _add_mat_suffix(matrix_path)
+
+        registrator.register(
+            target_image=fixed_image_path,
+            moving_image=moving_image_path,
+            output_image=transformed_image_path,
+            transform_file=matrix_path,
+            log_file=log_file_path,
+        )
+
+    def transform(
+        self,
+        fixed_image_path: str,
+        moving_image_path: str,
+        transformed_image_path: str,
+        matrix_path: str,
+        log_file_path: str = None,
+    ) -> None:
+        """
+        Apply a transformation using eReg.
+
+        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.
+            log_file_path (str): Path to the log file.
+        """
+        # TODO do we need to handle kwargs?
+        registrator = RegistrationClass(
+            configuration_file=self.configuration_file,
+        )
+
+        matrix_path = _add_mat_suffix(matrix_path)
+
+        registrator.resample_image(
+            target_image=fixed_image_path,
+            moving_image=moving_image_path,
+            output_image=transformed_image_path,
+            transform_file=matrix_path,
+            log_file=log_file_path,
+        )
+
+
+def _add_mat_suffix(filename: str) -> str:
+    """
+    Adds a ".mat" suffix to the filename if it doesn't have any extension.
+
+    Parameters:
+        filename (str): The filename to check and potentially modify.
+
+    Returns:
+        str: The filename with ".mat" suffix added if needed.
+    """
+    base, ext = os.path.splitext(filename)
+    if not ext:
+        filename += ".mat"
+    return filename

brainles_preprocessing/registration/niftyreg/niftyreg.py

@@ -1,4 +1,3 @@
-# TODO add typing and docs
 import os
 
 from auxiliary.runscript import ScriptRunner

example/example_modality_centric_preprocessor.py

@@ -6,7 +6,11 @@
 from brainles_preprocessing.brain_extraction import HDBetExtractor
 from brainles_preprocessing.modality import Modality
 from brainles_preprocessing.preprocessor import Preprocessor
-from brainles_preprocessing.registration import ANTsRegistrator, NiftyRegRegistrator
+from brainles_preprocessing.registration import (
+    ANTsRegistrator,
+    NiftyRegRegistrator,
+    eRegRegistrator,
+)
 
 
 def preprocess(inputDir):
@@ -103,6 +107,7 @@ def preprocess(inputDir):
             # choose the registration backend you want to use
             # registrator=NiftyRegRegistrator(),
             registrator=ANTsRegistrator(),
+            # registrator=eRegRegistrator(),
             brain_extractor=HDBetExtractor(),
             temp_folder="temporary_directory",
             limit_cuda_visible_devices="0",

pyproject.toml

@@ -59,10 +59,13 @@ rich = "^13.6.0"
 
 # optional registration backends
 antspyx = { version = "^0.4.2", optional = true }
+ereg = { version = "^0.0.10", optional = true }
+
 
 [tool.poetry.extras]
-all = ["antspyx"]
+all = ["antspyx", "ereg"]
 ants = ["antspyx"]
+ereg = ["ereg"]
 
 
 [tool.poetry.dev-dependencies]

tests/registrator_base.py

@@ -0,0 +1,90 @@
+import os
+from abc import abstractmethod
+
+from auxiliary.turbopath import turbopath
+import shutil
+
+
+class RegistratorBase:
+
+    @abstractmethod
+    def get_registrator(self):
+        pass
+
+    @abstractmethod
+    def get_method_and_extension(self):
+        pass
+
+    def setUp(self):
+        self.registrator = self.get_registrator()
+        self.method_name, self.matrix_extension = self.get_method_and_extension()
+
+        test_data_dir = turbopath(__file__).parent + "/test_data"
+        input_dir = test_data_dir + "/input"
+        self.output_dir = test_data_dir + f"/temp_output_{self.method_name}"
+        os.makedirs(self.output_dir, exist_ok=True)
+
+        self.fixed_image = input_dir + "/tcia_example_t1c.nii.gz"
+        self.moving_image = input_dir + "/tcia_example_t1.nii.gz"
+
+        self.matrix = self.output_dir + f"/{self.method_name}_matrix"
+        self.transform_matrix = input_dir + f"/{self.method_name}_matrix"
+
+    def tearDown(self):
+        # Clean up created files if they exist
+        shutil.rmtree(self.output_dir)
+        # pass
+
+    def test_register_creates_output_files(self):
+        transformed_image = (
+            self.output_dir + f"/{self.method_name}_registered_image.nii.gz"
+        )
+        log_file = self.output_dir + f"/{self.method_name}_registration.log"
+
+        self.registrator.register(
+            fixed_image_path=self.fixed_image,
+            moving_image_path=self.moving_image,
+            transformed_image_path=transformed_image,
+            matrix_path=self.matrix,
+            log_file_path=log_file,
+        )
+
+        self.assertTrue(
+            os.path.exists(transformed_image),
+            "transformed file was not created.",
+        )
+
+        self.assertTrue(
+            os.path.exists(f"{self.matrix}.{self.matrix_extension}"),
+            "matrix file was not created.",
+        )
+
+        self.assertTrue(
+            os.path.exists(log_file),
+            "log file was not created.",
+        )
+
+    def test_transform_creates_output_files(self):
+        transformed_image = (
+            self.output_dir + f"/{self.method_name}_transformed_image.nii.gz"
+        )
+        log_file = self.output_dir + f"/{self.method_name}_transformation.log"
+
+        print("tf m:", self.transform_matrix)
+        self.registrator.transform(
+            fixed_image_path=self.fixed_image,
+            moving_image_path=self.moving_image,
+            transformed_image_path=transformed_image,
+            matrix_path=self.transform_matrix,
+            log_file_path=log_file,
+        )
+
+        self.assertTrue(
+            os.path.exists(transformed_image),
+            "transformed file was not created.",
+        )
+
+        self.assertTrue(
+            os.path.exists(log_file),
+            "log file was not created.",
+        )

tests/test_data/input/niftyreg_matrix.txt

@@ -0,0 +1,4 @@
+0.9888756 0.1354215 0.06153017 -1.314461
+-0.1397797 0.9874795 0.07311541 -1.147088
+-0.05085838 -0.08090271 0.9954236 2.731441
+0 0 0 1

tests/test_registrators.py

@@ -1,93 +1,10 @@
-import os
-import shutil
-import unittest
-from abc import abstractmethod
-
-from auxiliary.turbopath import turbopath
+from registrator_base import RegistratorBase
 
 from brainles_preprocessing.registration.ANTs.ANTs import ANTsRegistrator
+from brainles_preprocessing.registration.eReg.eReg import eRegRegistrator
 from brainles_preprocessing.registration.niftyreg.niftyreg import NiftyRegRegistrator
 
-
-class RegistratorBase:
-    @abstractmethod
-    def get_registrator(self):
-        pass
-
-    @abstractmethod
-    def get_method_and_extension(self):
-        pass
-
-    def setUp(self):
-        self.registrator = self.get_registrator()
-        self.method_name, self.matrix_extension = self.get_method_and_extension()
-
-        test_data_dir = turbopath(__file__).parent + "/test_data"
-        input_dir = test_data_dir + "/input"
-        self.output_dir = test_data_dir + f"/temp_output_{self.method_name}"
-        os.makedirs(self.output_dir, exist_ok=True)
-
-        self.fixed_image = input_dir + "/tcia_example_t1c.nii.gz"
-        self.moving_image = input_dir + "/bet_tcia_example_t1c_mask.nii.gz"
-
-        self.matrix = self.output_dir + "/matrix"
-        self.transform_matrix = input_dir + f"/matrix.{self.matrix_extension}"
-
-    def tearDown(self):
-        # Clean up created files if they exist
-        shutil.rmtree(self.output_dir)
-
-    def test_register_creates_output_files(self):
-        transformed_image = self.output_dir + "/registered_image.nii.gz"
-        log_file = self.output_dir + "/registration.log"
-
-        self.registrator.register(
-            fixed_image_path=self.fixed_image,
-            moving_image_path=self.moving_image,
-            transformed_image_path=transformed_image,
-            matrix_path=self.matrix,
-            log_file_path=log_file,
-        )
-
-        self.assertTrue(
-            os.path.exists(transformed_image),
-            "transformed file was not created.",
-        )
-
-        self.assertTrue(
-            os.path.exists(f"{self.matrix}.{self.matrix_extension}"),
-            "matrix file was not created.",
-        )
-
-        self.assertTrue(
-            os.path.exists(log_file),
-            "log file was not created.",
-        )
-
-    def test_transform_creates_output_files(self):
-        transformed_image = self.output_dir + "/transformed_image.nii.gz"
-        log_file = self.output_dir + "/transformation.log"
-
-        self.registrator.transform(
-            fixed_image_path=self.fixed_image,
-            moving_image_path=self.moving_image,
-            transformed_image_path=transformed_image,
-            matrix_path=self.transform_matrix,
-            log_file_path=log_file,
-        )
-
-        self.assertTrue(
-            os.path.exists(transformed_image),
-            "transformed file was not created.",
-        )
-
-        self.assertTrue(
-            os.path.exists(log_file),
-            "log file was not created.",
-        )
-
-
-# TODO also test transform
+import unittest
 
 
 class TestANTsRegistrator(RegistratorBase, unittest.TestCase):
@@ -104,3 +21,11 @@ def get_registrator(self):
 
     def get_method_and_extension(self):
         return "niftyreg", "txt"
+
+
+class TestEregRegistratorRegistrator(RegistratorBase, unittest.TestCase):
+    def get_registrator(self):
+        return eRegRegistrator()
+
+    def get_method_and_extension(self):
+        return "ereg", "mat"