Merge pull request #136 from BrainLesion/135-feature-generalize-back-to-native-class-to-allow-transforms-in-both-directions

135 feature generalize back to native class to allow transforms in both directions

Author: neuronflow

README.md

@@ -10,9 +10,10 @@
 
 `BrainLes preprocessing` is a comprehensive tool for preprocessing tasks in biomedical imaging, with a focus on (but not limited to) multi-modal brain MRI. It can be used to build modular preprocessing pipelines:
 
-This includes **normalization**, **co-registration**, **atlas registration** and **skulstripping / brain extraction**.
+This includes **normalization**, **co-registration**, **atlas registration**, **skullstripping / brain extraction**, **N4 Bias correction** and **defacing**.
+We provide means to transform images and segmentations in both directions between native and atlas space.
 
-BrainLes is written `backend-agnostic` meaning it allows to swap the registration, brain extraction tools and defacing tools.
+BrainLes is written modular and `backend-agnostic` meaning it allows to skip or swap registration, brain extraction, N4 bias correction and defacing tools.
 
 <!-- TODO include image here -->
 
@@ -86,7 +87,7 @@ moving_modalities = [
     )
 ]
 
-# instantiate and run the preprocessor using defaults for registration/ brain extraction/ defacing backends
+# instantiate and run the preprocessor using defaults for backends (registration, brain extraction, bias correction, defacing)
 preprocessor = Preprocessor(
     center_modality=center,
     moving_modalities=moving_modalities,
@@ -127,6 +128,9 @@ We currently provide support for [ANTs](https://github.com/ANTsX/ANTs) (default)
 We provide the SRI-24 atlas from this [publication](https://www.ncbi.nlm.nih.gov/pmc/articles/PMC2915788/).
 However, custom atlases in NIfTI format are supported.
 
+### N4 Bias correction
+We currently provide support for N4 Bias correction based on [SimpleITK](https://simpleitk.org/)
+
 ### Brain extraction
 We currently provide support for [HD-BET](https://github.com/MIC-DKFZ/HD-BET).
 

brainles_preprocessing/registration/ANTs/ANTs.py

@@ -142,7 +142,7 @@ def transform(
         transformed_image_path: Union[str, Path],
         matrix_path: str | Path | List[str | Path],
         log_file_path: Union[str, Path],
-        interpolator: str = "linear",
+        interpolator: str = "nearestNeighbor",
         **kwargs,
     ) -> None:
         """
@@ -154,7 +154,7 @@ def transform(
             transformed_image_path (str or Path): Path to the transformed image (output).
             matrix_path (str or Path or List[str | Path]): Path to the transformation matrix or a list of matrices.
             log_file_path (str or Path): Path to the log file.
-            interpolator (str): Interpolator to use for the transformation. Default is 'linear'.
+            interpolator (str): Interpolator to use for the transformation. Default is 'nearestNeighbor'.
             **kwargs: Additional transformation parameters to update the instantiated defaults.
         Raises:
             AssertionError: If the interpolator is not valid.
@@ -231,7 +231,7 @@ def inverse_transform(
         transformed_image_path: Union[str, Path],
         matrix_path: str | Path | List[str | Path],
         log_file_path: Union[str, Path],
-        interpolator: str = "linear",
+        interpolator: str = "nearestNeighbor",
         **kwargs,
     ) -> None:
         """
@@ -243,7 +243,7 @@ def inverse_transform(
             transformed_image_path (str or Path): Path to the transformed image (output).
             matrix_path (str or Path): Path to the transformation matrix.
             log_file_path (str or Path): Path to the log file.
-            interpolator (str): Interpolator to use for the transformation. Default is 'linear'.
+            interpolator (str): Interpolator to use for the transformation. Default is 'nearestNeighbor'.
             **kwargs: Additional transformation parameters to update the instantiated defaults.
         """
         if not isinstance(matrix_path, list):

brainles_preprocessing/registration/niftyreg/niftyreg.py

@@ -150,7 +150,7 @@ def transform(
         transformed_image_path: str,
         matrix_path: str | Path | List[str | Path],
         log_file_path: str,
-        interpolator: str = "1",
+        interpolator: str = "0",
         **kwargs: dict,
     ) -> None:
         """
@@ -226,7 +226,7 @@ def inverse_transform(
         transformed_image_path: str,
         matrix_path: List[str | Path],
         log_file_path: str,
-        interpolator: str = "1",
+        interpolator: str = "0",
     ) -> None:
         """
         Apply inverse transformation using NiftyReg.

brainles_preprocessing/back_to_native.py renamed to brainles_preprocessing/transform.py

@@ -1,25 +1,32 @@
 from pathlib import Path
-from typing import List, Optional, Union
+from typing import Optional, Union
 
-from brainles_preprocessing.constants import Atlas, PreprocessorSteps
-from brainles_preprocessing.defacing import Defacer, QuickshearDefacer
-from brainles_preprocessing.modality import CenterModality, Modality
 from brainles_preprocessing.registration import ANTsRegistrator
 from brainles_preprocessing.registration.registrator import Registrator
 from brainles_preprocessing.utils.logging_utils import LoggingManager
-from brainles_preprocessing.utils.zenodo import verify_or_download_atlases
 
 logging_man = LoggingManager(name=__name__)
 logger = logging_man.get_logger()
 
 
-class BackToNativeSpace:
+class Transform:
+    """
+    Class to apply precomputed transformations based on the registration process.
+    Common use case is to apply inverse transformations to transform e.g. segmentations in atlas space back to native space or transform existing labels in native space to atlas space.
+    """
 
     def __init__(
         self,
         transformations_dir: Union[str, Path],
         registrator: Optional[Registrator] = None,
-    ):
+    ) -> None:
+        """
+        Initialize the Transform class.
+
+        Args:
+            transformations_dir (Union[str, Path]): Directory containing precomputed transformations for each modality.
+            registrator (Optional[Registrator], optional): Registrator instance to use for applying transformations. If None, defaults to ANTsRegistrator.
+        """
 
         self.transformations_dir = Path(transformations_dir)
 
@@ -29,17 +36,18 @@ def __init__(
             )
         self.registrator: Registrator = registrator or ANTsRegistrator()
 
-    def transform(
+    def apply(
         self,
         target_modality_name: str,
         target_modality_img: Union[str, Path],
         moving_image: Union[str, Path],
         output_img_path: Union[str, Path],
         log_file_path: Union[str, Path],
         interpolator: Optional[str] = None,
-    ):
+        inverse: bool = False,
+    ) -> None:
         """
-        Apply inverse transformation to a moving image to align it with a target modality.
+        Apply forward/ inverse transformation to a moving image to align it with a target modality.
 
         Args:
             target_modality_name (str): Name of the target modality. Must match the name used to create the transformations.
@@ -48,11 +56,12 @@ def transform(
             output_img_path (Union[str, Path]): Path where the transformed image will be saved.
             log_file_path (Union[str, Path]): Path to the log file where transformation details will be written.
             interpolator (Optional[str]): Interpolation method used during transformation.
+            inverse (bool): If True, applies the inverse transformation. Default is False.
                 Available options depend on the chosen registrator:
 
                 - **ANTsRegistrator**:
-                    - "linear" (default)
-                    - "nearestNeighbor"
+                    - "linear"
+                    - "nearestNeighbor" (default)
                     - "multiLabel" (deprecated, prefer "genericLabel")
                     - "gaussian"
                     - "bSpline"
@@ -63,16 +72,16 @@ def transform(
                     - "genericLabel" (recommended for label images)
 
                 - **NiftyReg**:
-                    - "0": nearest neighbor
-                    - "1": linear (default)
+                    - "0": nearest neighbor (default)
+                    - "1": linear
                     - "3": cubic spline
                     - "4": sinc
 
         Raises:
             AssertionError: If the transformations directory for the given modality does not exist.
         """
         logger.info(
-            f"Applying inverse transformation for {target_modality_name} using {self.registrator.__class__.__name__}."
+            f"Applying {'inverse' if inverse else ''} transformation for {target_modality_name} using {self.registrator.__class__.__name__}."
         )
 
         # assert modality name eixsts in transformations_dir
@@ -86,7 +95,9 @@ def transform(
 
         transforms = list(modality_transformations_dir.iterdir())
         transforms.sort()  # sort by name to get order for forward transform
-        transforms = transforms[::-1]  # inverse order for inverse transform
+
+        if inverse:
+            transforms = transforms[::-1]  # inverse order for inverse transform
 
         kwargs = {
             "fixed_image_path": target_modality_img,
@@ -97,4 +108,8 @@ def transform(
         }
         if interpolator is not None:
             kwargs["interpolator"] = interpolator
-        self.registrator.inverse_transform(**kwargs)
+
+        if inverse:
+            self.registrator.inverse_transform(**kwargs)
+        else:
+            self.registrator.transform(**kwargs)

tests/test_preprocessor.py

@@ -0,0 +1,93 @@
+from types import SimpleNamespace
+
+import pytest
+
+from brainles_preprocessing.preprocessor import CenterModality, Modality, Preprocessor
+
+
+# --- Dummy classes for testing ---
+class DummyModality:
+    def __init__(self, name):
+        self.modality_name = name
+
+
+class DummyCenterModality(DummyModality):
+    pass
+
+
+@pytest.fixture
+def dummy_registrator():
+    return SimpleNamespace()
+
+
+@pytest.fixture
+def dummy_brain_extractor():
+    return SimpleNamespace()
+
+
+@pytest.fixture
+def dummy_defacer():
+    return SimpleNamespace()
+
+
+# --- Tests ---
+
+
+def test_no_name_conflicts(dummy_registrator, dummy_brain_extractor, dummy_defacer):
+    center = CenterModality("T1", input_path="", raw_skull_output_path="tmp")
+    moving = [
+        Modality("T2", input_path="", raw_skull_output_path="tmp"),
+        Modality("FLAIR", input_path="", raw_skull_output_path="tmp"),
+    ]
+    # Should not raise
+    Preprocessor(
+        center_modality=center,
+        moving_modalities=moving,
+        registrator=dummy_registrator,
+        brain_extractor=dummy_brain_extractor,
+        defacer=dummy_defacer,
+    )
+
+
+def test_single_duplicate_name_raises(
+    dummy_registrator, dummy_brain_extractor, dummy_defacer
+):
+    center = CenterModality("T1", input_path="", raw_skull_output_path="tmp")
+    moving = [
+        Modality("T1", input_path="", raw_skull_output_path="tmp"),  # Duplicate name
+        Modality("FLAIR", input_path="", raw_skull_output_path="tmp"),
+    ]
+
+    with pytest.raises(ValueError, match=r"Duplicate modality names found: T1"):
+        Preprocessor(
+            center_modality=center,
+            moving_modalities=moving,
+            registrator=dummy_registrator,
+            brain_extractor=dummy_brain_extractor,
+            defacer=dummy_defacer,
+        )
+
+
+def test_multiple_duplicate_names_raises(
+    dummy_registrator, dummy_brain_extractor, dummy_defacer
+):
+    center = CenterModality("T1", input_path="", raw_skull_output_path="tmp")
+    moving = [
+        Modality("T1", input_path="", raw_skull_output_path="tmp"),  # Duplicate
+        Modality("FLAIR", input_path="", raw_skull_output_path="tmp"),
+        Modality("FLAIR", input_path="", raw_skull_output_path="tmp"),  # Duplicate
+    ]
+
+    with pytest.raises(ValueError) as exc_info:
+        Preprocessor(
+            center_modality=center,
+            moving_modalities=moving,
+            registrator=dummy_registrator,
+            brain_extractor=dummy_brain_extractor,
+            defacer=dummy_defacer,
+        )
+
+    # Check that all duplicates are reported in the error message
+    error_msg = str(exc_info.value)
+    assert "T1" in error_msg
+    assert "FLAIR" in error_msg