Merge pull request #131 from BrainLesion/92-unite-registration-steps-into-one-transformation

92 unite registration steps into one transformation

Author: neuronflow

brainles_preprocessing/back_to_native.py

@@ -0,0 +1,100 @@
+from pathlib import Path
+from typing import List, 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:
+
+    def __init__(
+        self,
+        transformations_dir: Union[str, Path],
+        registrator: Optional[Registrator] = None,
+    ):
+
+        self.transformations_dir = Path(transformations_dir)
+
+        if registrator is None:
+            logger.warning(
+                "No registrator provided, using default ANTsRegistrator for registration."
+            )
+        self.registrator: Registrator = registrator or ANTsRegistrator()
+
+    def transform(
+        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,
+    ):
+        """
+        Apply 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.
+            target_modality_img (Union[str, Path]): Path to the target modality image.
+            moving_image (Union[str, Path]): Path to the moving image. E.g., this could be a segmentation in atlas space.
+            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.
+                Available options depend on the chosen registrator:
+
+                - **ANTsRegistrator**:
+                    - "linear" (default)
+                    - "nearestNeighbor"
+                    - "multiLabel" (deprecated, prefer "genericLabel")
+                    - "gaussian"
+                    - "bSpline"
+                    - "cosineWindowedSinc"
+                    - "welchWindowedSinc"
+                    - "hammingWindowedSinc"
+                    - "lanczosWindowedSinc"
+                    - "genericLabel" (recommended for label images)
+
+                - **NiftyReg**:
+                    - "0": nearest neighbor
+                    - "1": linear (default)
+                    - "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__}."
+        )
+
+        # assert modality name eixsts in transformations_dir
+        modality_transformations_dir = (
+            self.transformations_dir / f"{target_modality_name}"
+        )
+
+        assert (
+            modality_transformations_dir.exists()
+        ), f"Transformations directory for {target_modality_name} does not exist: {modality_transformations_dir}"
+
+        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
+
+        kwargs = {
+            "fixed_image_path": target_modality_img,
+            "moving_image_path": moving_image,
+            "transformed_image_path": output_img_path,
+            "matrix_path": transforms,
+            "log_file_path": str(log_file_path),
+        }
+        if interpolator is not None:
+            kwargs["interpolator"] = interpolator
+        self.registrator.inverse_transform(**kwargs)

brainles_preprocessing/modality.py

@@ -2,13 +2,18 @@
 import shutil
 import warnings
 from pathlib import Path
-from typing import Optional, Union
+from typing import Dict, Optional, Union
 
 from auxiliary.nifti.io import read_nifti, write_nifti
+
 from brainles_preprocessing.brain_extraction.brain_extractor import BrainExtractor
 from brainles_preprocessing.constants import PreprocessorSteps
 from brainles_preprocessing.defacing import Defacer, QuickshearDefacer
 from brainles_preprocessing.normalization.normalizer_base import Normalizer
+from brainles_preprocessing.registration import (  # TODO: this will throw warnings if ANTs or NiftyReg are not installed, not ideal
+    ANTsRegistrator,
+    NiftyRegRegistrator,
+)
 from brainles_preprocessing.registration.registrator import Registrator
 
 logger = logging.getLogger(__name__)
@@ -42,6 +47,7 @@ class Modality:
         normalized_defaced_output_path (str or Path, optional): Path to save the normalized defaced modality data. Requires a normalizer.
         bet (bool): Indicates whether brain extraction is enabled.
         atlas_correction (bool): Indicates whether atlas correction should be performed.
+        coregistration_transform_path (str or None): Path to the coregistration transformation matrix, will be set after coregistration.
 
     Example:
         >>> t1_modality = Modality(
@@ -73,6 +79,7 @@ def __init__(
         self.current = self.input_path
         self.normalizer = normalizer
         self.atlas_correction = atlas_correction
+        self.transformation_paths: Dict[PreprocessorSteps, Path | None] = {}
 
         # Check that atleast one output is generated
         if not any(
@@ -128,6 +135,7 @@ def __init__(
             self.normalized_defaced_output_path = None
 
         self.steps = {k: None for k in PreprocessorSteps}
+        self.steps[PreprocessorSteps.INPUT] = self.input_path
 
     @property
     def bet(self) -> bool:
@@ -197,6 +205,26 @@ def normalize(
         else:
             logger.info("No normalizer specified; skipping normalization.")
 
+    def _find_transformation_matrix(
+        self, transform_incomplete_path: Path
+    ) -> Optional[Path]:
+        possible_Files = list(
+            transform_incomplete_path.parent.glob(f"{transform_incomplete_path.stem}.*")
+        )
+        if len(possible_Files) == 0:
+            logger.warning(
+                f"No transformation matrix found for {transform_incomplete_path}. "
+                "Returning None."
+            )
+            return None
+        elif len(possible_Files) > 1:
+            # TODO: Handle this case more gracefully, e.g., try to find proper extension based on the registrator
+            logger.warning(
+                f"Multiple transformation matrices found for {transform_incomplete_path}. "
+                "Returning the first one."
+            )
+        return possible_Files[0]
+
     def register(
         self,
         registrator: Registrator,
@@ -213,6 +241,7 @@ def register(
             fixed_image_path (str or Path): Path to the fixed image.
             registration_dir (str or Path): Directory to store registration results.
             moving_image_name (str): Name of the moving image.
+            step (PreprocessorSteps): The current preprocessing step.
 
         Returns:
             Path: Path to the registration matrix.
@@ -224,7 +253,7 @@ def register(
         registered_log = registration_dir / f"{moving_image_name}.log"
 
         # Note, add file ending depending on registration backend!
-        registered_matrix = registration_dir / f"{moving_image_name}"
+        registered_matrix = registration_dir / f"M_{moving_image_name}"
 
         registrator.register(
             fixed_image_path=fixed_image_path,
@@ -235,6 +264,11 @@ def register(
         )
         self.current = registered
         self.steps[step] = registered
+
+        self.transformation_paths[step] = self._find_transformation_matrix(
+            transform_incomplete_path=registered_matrix
+        )
+
         return registered_matrix
 
     def apply_bet_mask(
@@ -320,7 +354,7 @@ def transform(
             registration_dir_path (str or Path): Directory to store transformation results.
             moving_image_name (str): Name of the moving image.
             transformation_matrix_path (str or Path): Path to the transformation matrix.
-
+            step (PreprocessorSteps): The current preprocessing step.
         Returns:
             None
         """
@@ -331,15 +365,42 @@ def transform(
         transformed = registration_dir_path / f"{moving_image_name}.nii.gz"
         transformed_log = registration_dir_path / f"{moving_image_name}.log"
 
-        registrator.transform(
-            fixed_image_path=fixed_image_path,
-            moving_image_path=self.current,
-            transformed_image_path=transformed,
-            matrix_path=transformation_matrix_path,
-            log_file_path=transformed_log,
-        )
+        if (
+            isinstance(registrator, (ANTsRegistrator, NiftyRegRegistrator))
+            and step == PreprocessorSteps.ATLAS_REGISTERED
+        ):
+            # we test uniting transforms for these registrators
+            assert (
+                self.transformation_paths.get(PreprocessorSteps.COREGISTERED, None)
+                is not None
+            ), "Coregistration must be performed before applying atlas registration."
+
+            registrator.transform(
+                fixed_image_path=fixed_image_path,
+                moving_image_path=self.steps[PreprocessorSteps.INPUT],
+                transformed_image_path=transformed,
+                matrix_path=[
+                    self.transformation_paths[
+                        PreprocessorSteps.COREGISTERED
+                    ],  # coregistration matrix
+                    transformation_matrix_path,  # atlas registration matrix
+                ],
+                log_file_path=transformed_log,
+            )
+        else:
+            registrator.transform(
+                fixed_image_path=fixed_image_path,
+                moving_image_path=self.current,
+                transformed_image_path=transformed,
+                matrix_path=transformation_matrix_path,
+                log_file_path=transformed_log,
+            )
+
         self.current = transformed
         self.steps[step] = transformed
+        self.transformation_paths[step] = self._find_transformation_matrix(
+            transform_incomplete_path=transformation_matrix_path
+        )
 
     def extract_brain_region(
         self,

brainles_preprocessing/preprocessor.py

@@ -164,6 +164,7 @@ def run(
         save_dir_atlas_correction: Optional[Union[str, Path]] = None,
         save_dir_brain_extraction: Optional[Union[str, Path]] = None,
         save_dir_defacing: Optional[Union[str, Path]] = None,
+        save_dir_transformations: Optional[Union[str, Path]] = None,
         log_file: Optional[Union[str, Path]] = None,
     ):
         """
@@ -176,6 +177,7 @@ def run(
             save_dir_atlas_correction (str or Path, optional): Directory path to save intermediate atlas correction results.
             save_dir_brain_extraction (str or Path, optional): Directory path to save intermediate brain extraction results.
             save_dir_defacing (str or Path, optional): Directory path to save intermediate defacing results.
+            save_dir_transformations (str or Path, optional): Directory path to save transformation matrices. Defaults to None.
             log_file (str or Path, optional): Path to save the log file. Defaults to a timestamped file in the current directory.
 
         This method orchestrates the entire preprocessing workflow by sequentially performing:
@@ -250,6 +252,28 @@ def run(
             save_dir_defacing=save_dir_defacing,
         )
 
+        # move to separate method
+        if save_dir_transformations:
+            save_dir_transformations = Path(save_dir_transformations)
+
+            # Save transformation matrices
+            logger.info(f"Saving transformation matrices to {save_dir_transformations}")
+            for modality in self.all_modalities:
+
+                modality_transformations_dir = (
+                    save_dir_transformations / modality.modality_name
+                )
+                modality_transformations_dir.mkdir(exist_ok=True, parents=True)
+                for step, path in modality.transformation_paths.items():
+                    if path is not None:
+                        shutil.copyfile(
+                            src=str(path.absolute()),
+                            dst=str(
+                                modality_transformations_dir
+                                / f"{step.value}_{path.name}"
+                            ),
+                        )
+
         # End
         logger.info(f"{' Preprocessing complete ':=^80}")