Merge pull request #93 from BrainLesion/88-feature-request-allow-user…

…-to-save-bet-and-defacing-masks

88 feature request allow user to save bet and defacing masks

README.md

@@ -12,9 +12,7 @@
 
 This includes **normalization**, **co-registration**, **atlas registration** and **skulstripping / brain extraction**.
 
-BrainLes is written `backend-agnostic` meaning it allows to swap the registration and brain extraction tools.
-
-<!-- TODO mention defacing -->
+BrainLes is written `backend-agnostic` meaning it allows to swap the registration, brain extraction tools and defacing tools.
 
 <!-- TODO include image here -->
 
@@ -32,7 +30,7 @@ pip install brainles-preprocessing
 A minimal example to register (to the standard atlas using ANTs) and skull strip (using HDBet) a t1c image (center modality) with 1 moving modality (flair) could look like this:
 ```python
 from pathlib import Path
-from brainles_preprocessing.modality import Modality
+from brainles_preprocessing.modality import Modality, CenterModality
 from brainles_preprocessing.normalization.percentile_normalizer import (
     PercentileNormalizer,
 )
@@ -48,8 +46,8 @@ percentile_normalizer = PercentileNormalizer(
     upper_limit=1,
 )
 
-# define modalities
-center = Modality(
+# define center and moving modalities
+center = CenterModality(
     modality_name="t1c",
     input_path=patient_folder / "t1c.nii.gz",
     normalizer=percentile_normalizer,
@@ -58,6 +56,9 @@ center = Modality(
     raw_bet_output_path="patient/raw_bet_dir/t1c_bet_raw.nii.gz",
     normalized_skull_output_path="patient/norm_skull_dir/t1c_skull_normalized.nii.gz",
     normalized_bet_output_path="patient/norm_bet_dir/t1c_bet_normalized.nii.gz",
+    # specify output paths for the brain extraction and defacing masks
+    bet_mask_output_path="patient/masks/t1c_bet_mask.nii.gz",
+    defacing_mask_output_path="patient/masks/t1c_defacing_mask.nii.gz",
 )
 
 moving_modalities = [
@@ -107,15 +108,15 @@ We provide a (WIP) documentation. Have a look [here](https://brainles-preprocess
 ## FAQ
 Please credit the authors by citing their work.
 
+### Registration
+We currently provide support for [ANTs](https://github.com/ANTsX/ANTs) (default), [Niftyreg](https://github.com/KCL-BMEIS/niftyreg) (Linux), eReg (experimental)
+
 ### Atlas Reference
 We provide the SRI-24 atlas from this [publication](https://www.ncbi.nlm.nih.gov/pmc/articles/PMC2915788/).
-However, custom atlases can be supplied.
+However, custom atlases in NIfTI format are supported.
 
 ### Brain extraction
 We currently provide support for [HD-BET](https://github.com/MIC-DKFZ/HD-BET).
 
-### Registration
-We currently provide support for [ANTs](https://github.com/ANTsX/ANTs) (default), [Niftyreg](https://github.com/KCL-BMEIS/niftyreg) (Linux), eReg (experimental)
-
 ### Defacing
-We currently provide support for [Quickshear](https://github.com/nipy/quickshear)
+We currently provide support for [Quickshear](https://github.com/nipy/quickshear).

brainles_preprocessing/modality.py

@@ -1,5 +1,6 @@
 import logging
 import shutil
+import warnings
 from pathlib import Path
 from typing import Optional, Union
 
@@ -346,6 +347,8 @@ def extract_brain_region(
         bet_dir_path: Union[str, Path],
     ) -> Path:
         """
+        WARNING: Legacy method. Please Migrate to use the CenterModality Class. Will be removed in future versions.
+
         Extract the brain region using the specified brain extractor.
 
         Args:
@@ -355,6 +358,12 @@ def extract_brain_region(
         Returns:
             Path: Path to the extracted brain mask.
         """
+
+        warnings.warn(
+            "Legacy method. Please Migrate to use the CenterModality Class. Will be removed in future versions.",
+            category=DeprecationWarning,
+        )
+
         bet_dir_path = Path(bet_dir_path)
         bet_log = bet_dir_path / "brain-extraction.log"
 
@@ -382,6 +391,8 @@ def deface(
         defaced_dir_path: Union[str, Path],
     ) -> Path:
         """
+        WARNING: Legacy method. Please Migrate to use the CenterModality Class. Will be removed in future versions.
+
         Deface the current modality using the specified defacer.
 
         Args:
@@ -391,7 +402,10 @@ def deface(
         Returns:
             Path: Path to the extracted brain mask.
         """
-
+        warnings.warn(
+            "Legacy method. Please Migrate to use the CenterModality class. Will be removed in future versions.",
+            category=DeprecationWarning,
+        )
         if isinstance(defacer, QuickshearDefacer):
 
             defaced_dir_path = Path(defaced_dir_path)
@@ -443,3 +457,179 @@ def save_current_image(
                 src=str(self.current),
                 dst=str(output_path),
             )
+
+
+class CenterModality(Modality):
+    """
+    Represents a medical image center modality with associated preprocessing information.
+
+    Args:
+        modality_name (str): Name of the modality, e.g., "T1", "T2", "FLAIR".
+        input_path (str or Path): Path to the input modality data.
+        normalizer (Normalizer, optional): An optional normalizer for intensity normalization.
+        raw_bet_output_path (str or Path, optional): Path to save the raw brain extracted modality data.
+        raw_skull_output_path (str or Path, optional): Path to save the raw modality data with skull.
+        raw_defaced_output_path (str or Path, optional): Path to save the raw defaced modality data.
+        normalized_bet_output_path (str or Path, optional): Path to save the normalized brain extracted modality data. Requires a normalizer.
+        normalized_skull_output_path (str or Path, optional): Path to save the normalized modality data with skull. Requires a normalizer.
+        normalized_defaced_output_path (str or Path, optional): Path to save the normalized defaced modality data. Requires a normalizer.
+        atlas_correction (bool, optional): Indicates whether atlas correction should be performed.
+        bet_mask_output_path (str or Path, optional): Path to save the brain extraction mask.
+        defacing_mask_output_path (str or Path, optional): Path to save the defacing mask.
+
+    Attributes:
+        modality_name (str): Name of the modality.
+        input_path (str or Path): Path to the input modality data.
+        normalizer (Normalizer, optional): An optional normalizer for intensity normalization.
+        raw_bet_output_path (str or Path, optional): Path to save the raw brain extracted modality data.
+        raw_skull_output_path (str or Path, optional): Path to save the raw modality data with skull.
+        raw_defaced_output_path (str or Path, optional): Path to save the raw defaced modality data.
+        normalized_bet_output_path (str or Path, optional): Path to save the normalized brain extracted modality data. Requires a normalizer.
+        normalized_skull_output_path (str or Path, optional): Path to save the normalized modality data with skull. Requires a normalizer.
+        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.
+        bet_mask_output_path (Path, optional): Path to save the brain extraction mask.
+        defacing_mask_output_path (Path, optional): Path to save the defacing mask.
+
+    Example:
+        >>> t1_modality = CenterModality(
+        ...     modality_name="T1",
+        ...     input_path="/path/to/input_t1.nii",
+        ...     normalizer=PercentileNormalizer(),
+        ...     raw_bet_output_path="/path/to/raw_bet_t1.nii",
+        ...     normalized_bet_output_path="/path/to/norm_bet_t1.nii",
+        ...     bet_mask_output_path="/path/to/bet_mask_t1.nii",
+        ... )
+    """
+
+    def __init__(
+        self,
+        modality_name: str,
+        input_path: Union[str, Path],
+        normalizer: Optional[Normalizer] = None,
+        raw_bet_output_path: Optional[Union[str, Path]] = None,
+        raw_skull_output_path: Optional[Union[str, Path]] = None,
+        raw_defaced_output_path: Optional[Union[str, Path]] = None,
+        normalized_bet_output_path: Optional[Union[str, Path]] = None,
+        normalized_skull_output_path: Optional[Union[str, Path]] = None,
+        normalized_defaced_output_path: Optional[Union[str, Path]] = None,
+        atlas_correction: bool = True,
+        bet_mask_output_path: Optional[Union[str, Path]] = None,
+        defacing_mask_output_path: Optional[Union[str, Path]] = None,
+    ) -> None:
+        super().__init__(
+            modality_name=modality_name,
+            input_path=input_path,
+            normalizer=normalizer,
+            raw_bet_output_path=raw_bet_output_path,
+            raw_skull_output_path=raw_skull_output_path,
+            raw_defaced_output_path=raw_defaced_output_path,
+            normalized_bet_output_path=normalized_bet_output_path,
+            normalized_skull_output_path=normalized_skull_output_path,
+            normalized_defaced_output_path=normalized_defaced_output_path,
+            atlas_correction=atlas_correction,
+        )
+        # Only for CenterModality
+        self.bet_mask_output_path = (
+            Path(bet_mask_output_path) if bet_mask_output_path else None
+        )
+        self.defacing_mask_output_path = (
+            Path(defacing_mask_output_path) if defacing_mask_output_path else None
+        )
+
+    def extract_brain_region(
+        self,
+        brain_extractor: BrainExtractor,
+        bet_dir_path: Union[str, Path],
+    ) -> Path:
+        """
+        Extract the brain region using the specified brain extractor.
+
+        Args:
+            brain_extractor (BrainExtractor): The brain extractor object.
+            bet_dir_path (str or Path): Directory to store brain extraction results.
+
+        Returns:
+            Path: Path to the extracted brain mask.
+        """
+        bet_dir_path = Path(bet_dir_path)
+        bet_log = bet_dir_path / "brain-extraction.log"
+
+        atlas_bet_cm = bet_dir_path / f"atlas__{self.modality_name}_bet.nii.gz"
+        mask_path = bet_dir_path / f"atlas__{self.modality_name}_brain_mask.nii.gz"
+
+        brain_extractor.extract(
+            input_image_path=self.current,
+            masked_image_path=atlas_bet_cm,
+            brain_mask_path=mask_path,
+            log_file_path=bet_log,
+        )
+
+        if self.bet_mask_output_path:
+            logger.debug(f"Saving bet mask to {self.bet_mask_output_path}")
+            self.save_mask(mask_path=mask_path, output_path=self.bet_mask_output_path)
+
+        # always temporarily store bet image for center modality, since e.g. quickshear defacing could require it
+        # down the line even if the user does not wish to save the bet image
+        self.steps[PreprocessorSteps.BET] = atlas_bet_cm
+
+        if self.bet:
+            self.current = atlas_bet_cm
+        return mask_path
+
+    def deface(
+        self,
+        defacer,
+        defaced_dir_path: Union[str, Path],
+    ) -> Path:
+        """
+        Deface the current modality using the specified defacer.
+
+        Args:
+            defacer (Defacer): The defacer object.
+            defaced_dir_path (str or Path): Directory to store defacing results.
+
+        Returns:
+            Path: Path to the extracted brain mask.
+        """
+
+        if isinstance(defacer, QuickshearDefacer):
+
+            defaced_dir_path = Path(defaced_dir_path)
+            atlas_mask_path = (
+                defaced_dir_path / f"atlas__{self.modality_name}_deface_mask.nii.gz"
+            )
+
+            defacer.deface(
+                mask_image_path=atlas_mask_path,
+                input_image_path=self.steps[PreprocessorSteps.BET],
+            )
+
+            if self.defacing_mask_output_path:
+                logger.debug(f"Saving deface mask to {self.defacing_mask_output_path}")
+                self.save_mask(
+                    mask_path=atlas_mask_path,
+                    output_path=self.defacing_mask_output_path,
+                )
+
+            return atlas_mask_path
+        else:
+            logger.warning(
+                "Defacing method not implemented yet. Skipping defacing for this modality."
+            )
+            return None
+
+    def save_mask(self, mask_path: Union[str, Path], output_path: Path) -> None:
+        """
+        Save the mask to the specified output path.
+
+        Args:
+            mask_path (Union[str, Path]): Mask NifTI file path.
+            output_path (Path): Output NifTI file path.
+        """
+        output_path.parent.mkdir(parents=True, exist_ok=True)
+        shutil.copyfile(
+            src=str(mask_path),
+            dst=str(output_path),
+        )

brainles_preprocessing/normalization/__init__.py

@@ -0,0 +1,3 @@
+from .normalizer_base import Normalizer
+from .percentile_normalizer import PercentileNormalizer
+from .windowing_normalizer import WindowingNormalizer

brainles_preprocessing/preprocessor.py

@@ -10,12 +10,13 @@
 from functools import wraps
 from pathlib import Path
 from typing import List, Optional, Union
+import warnings
 
 from brainles_preprocessing.constants import PreprocessorSteps
 from brainles_preprocessing.defacing import Defacer, QuickshearDefacer
 
 from .brain_extraction.brain_extractor import BrainExtractor, HDBetExtractor
-from .modality import Modality
+from .modality import Modality, CenterModality
 from .registration import ANTsRegistrator
 from .registration.registrator import Registrator
 
@@ -30,7 +31,7 @@ class Preprocessor:
     Preprocesses medical image modalities using coregistration, normalization, brain extraction, and more.
 
     Args:
-        center_modality (Modality): The central modality for coregistration.
+        center_modality (CenterModality): The central modality for coregistration.
         moving_modalities (List[Modality]): List of modalities to be coregistered to the central modality.
         registrator (Registrator): The registrator object for coregistration and registration to the atlas.
         brain_extractor (Optional[BrainExtractor]): The brain extractor object for brain extraction.
@@ -44,7 +45,7 @@ class Preprocessor:
 
     def __init__(
         self,
-        center_modality: Modality,
+        center_modality: CenterModality,
         moving_modalities: List[Modality],
         registrator: Registrator = None,
         brain_extractor: Optional[BrainExtractor] = None,
@@ -56,6 +57,13 @@ def __init__(
     ):
         logging_man._setup_logger()
 
+        if not isinstance(center_modality, CenterModality):
+            warnings.warn(
+                "Center modality should be of type CenterModality instead of Modality to allow for more features, e.g. saving bet and deface masks. "
+                "Support for using Modality for the Center Modality will be deprecated in future versions. "
+                "Note: Moving modalities should still be of type Modality.",
+                category=DeprecationWarning,
+            )
         self.center_modality = center_modality
         self.moving_modalities = moving_modalities