Merge pull request #53 from BrainLesion/52-feature-request-clean-the-…

…normalization-and-brain-extraction-in-the-run-method

52-feature-request-clean-the-normalization-and-brain-extraction-in-the-run-method

brainles_preprocessing/modality.py

@@ -1,5 +1,6 @@
 import os
 import shutil
+from typing import List, Optional
 
 from auxiliary.nifti.io import read_nifti, write_nifti
 from auxiliary.normalization.normalizer_base import Normalizer
@@ -44,7 +45,7 @@ def __init__(
         output_path: str,
         bet: bool,
         atlas_correction: bool = True,
-        normalizer: Normalizer | None = None,
+        normalizer: Optional[Normalizer] = None,
     ) -> None:
         self.modality_name = modality_name
         self.input_path = turbopath(input_path)
@@ -221,5 +222,6 @@ def extract_brain_region(
             brain_mask_path=atlas_mask_path,
             log_file_path=bet_log,
         )
-        self.current = atlas_bet_cm
+        if self.bet is True:
+            self.current = atlas_bet_cm
         return atlas_mask_path

brainles_preprocessing/preprocessor.py

@@ -51,26 +51,27 @@ def __init__(
         self.atlas_dir = os.path.join(self.temp_folder, "atlas-space")
         os.makedirs(self.atlas_dir, exist_ok=True)
 
+    @property
+    def all_modalities(self):
+        return [self.center_modality] + self.moving_modalities
+
     def run(
         self,
-        brain_extraction: bool,  # TODO probably this should be true if one of the modalities has a bet flag with true?
-        normalization: bool,  # TODO probably this should be true if one of the modalities has a normalizer?
         save_dir_coregistration: Optional[str] = None,
         save_dir_atlas_registration: Optional[str] = None,
         save_dir_atlas_correction: Optional[str] = None,
         save_dir_brain_extraction: Optional[str] = None,
         save_dir_unnormalized: Optional[str] = None,
     ):
         """
-        Run the preprocessing pipeline.
+        Execute the preprocessing pipeline.
 
         Args:
-            brain_extraction (bool): Whether to perform brain extraction.
-            normalization (bool): Whether to perform intensity normalization.
-            save_dir_coregistration (str, optional): Directory to save coregistration results.
-            save_dir_atlas_registration (str, optional): Directory to save atlas registration results.
-            save_dir_brain_extraction (str, optional): Directory to save brain extraction results.
-            save_dir_unnormalized (str, optional): Directory to save unnormalized images.
+            save_dir_coregistration (str, optional): Directory path to save coregistration results.
+            save_dir_atlas_registration (str, optional): Directory path to save atlas registration results.
+            save_dir_atlas_correction (str, optional): Directory path to save atlas correction results.
+            save_dir_brain_extraction (str, optional): Directory path to save brain extraction results.
+            save_dir_unnormalized (str, optional): Directory path to save unnormalized images.
         """
         # Coregister moving modalities to center modality
         coregistration_dir = os.path.join(self.temp_folder, "coregistration")
@@ -84,9 +85,18 @@ def run(
                 registration_dir=coregistration_dir,
                 moving_image_name=file_name,
             )
-        self._save_coregistration(
-            coregistration_dir=coregistration_dir,
-            save_dir_coregistration=save_dir_coregistration,
+
+        shutil.copyfile(
+            src=self.center_modality.input_path,
+            dst=os.path.join(
+                coregistration_dir,
+                f"native__{self.center_modality.modality_name}.nii.gz",
+            ),
+        )
+
+        self._save_output(
+            src=coregistration_dir,
+            save_dir=save_dir_coregistration,
         )
 
         # Register center modality to atlas
@@ -142,6 +152,7 @@ def run(
         )
 
         # Optional: Brain extraction
+        brain_extraction = any(modality.bet for modality in self.all_modalities)
         if brain_extraction:
             bet_dir = os.path.join(self.temp_folder, "brain-extraction")
             os.makedirs(bet_dir, exist_ok=True)
@@ -164,6 +175,7 @@ def run(
             )
 
         # Optional: Normalization
+        normalization = any(modality.normalizer for modality in self.all_modalities)
         if normalization:
             for modality in [self.center_modality] + self.moving_modalities:
                 modality.normalize(
@@ -178,10 +190,6 @@ def run(
                 modality.output_path,
             )
 
-    @property
-    def all_modalities(self):
-        return [self.center_modality] + self.moving_modalities
-
     def _save_output(
         self,
         src: str,
@@ -195,28 +203,6 @@ def _save_output(
                 dirs_exist_ok=True,
             )
 
-    def _save_coregistration(
-        self,
-        coregistration_dir: str,
-        save_dir_coregistration: Optional[str],
-    ):
-        if save_dir_coregistration:
-            save_dir_coregistration = turbopath(save_dir_coregistration)
-            native_cm = os.path.join(
-                coregistration_dir,
-                f"native__{self.center_modality.modality_name}.nii.gz",
-            )
-
-            shutil.copyfile(
-                src=self.center_modality.input_path,
-                dst=native_cm,
-            )
-            shutil.copytree(
-                src=coregistration_dir,
-                dst=save_dir_coregistration,
-                dirs_exist_ok=True,
-            )
-
 
 class PreprocessorGPU(Preprocessor):
     """

example_modality_centric_preprocessor.py

@@ -97,8 +97,6 @@ def preprocess(inputDir):
             )
 
             preprocessor.run(
-                brain_extraction=True,
-                normalization=True,
                 save_dir_coregistration=brainles_dir + "/co-registration",
                 save_dir_atlas_registration=brainles_dir + "/atlas-registration",
                 save_dir_atlas_correction=brainles_dir + "/atlas-correction",