Changed workflow for setting the masking value in the skull stripping

brainles_preprocessing/brain_extraction/brain_extractor.py

@@ -1,12 +1,29 @@
 # TODO add typing and docs
 from abc import abstractmethod, ABC
 from pathlib import Path
-from typing import Union
+from typing import Optional, Union
+import numpy as np
 
 from auxiliary.io import read_image, write_image
 
 
 class BrainExtractor(ABC):
+    def __init__(
+        self,
+        masking_value: Optional[Union[int, float]] = None,
+    ):
+        """
+        Base class for skull stripping medical images using brain masks.
+
+        Subclasses should implement the `extract` method to generate a skull stripped image
+        based on the provided input image and mask.
+        """
+        # Just as in the defacer, masking value is a global value defined across all images and modalities
+        # If no value is passed, the minimum of a given input image is chosen
+        # TODO: Consider extending this to modality-specific masking values in the future, this should
+        # probably be implemented as a property of the specific modality
+        self.masking_value = masking_value
+
     @abstractmethod
     def extract(
         self,
@@ -22,7 +39,6 @@ def extract(
             input_image_path (str or Path): Path to the input image.
             masked_image_path (str or Path): Path where the brain-extracted image will be saved.
             brain_mask_path (str or Path): Path where the brain mask will be saved.
-            mode (str or Mode): Extraction mode.
             **kwargs: Additional keyword arguments.
         """
         pass
@@ -51,12 +67,8 @@ def apply_mask(
         except Exception as e:
             raise RuntimeError(f"Error reading files: {e}") from e
 
-        # Check that the input and mask have the same shape
-        if input_data.shape != mask_data.shape:
-            raise ValueError("Input image and mask must have the same dimensions.")
-
-        # Mask and save it
-        masked_data = input_data * mask_data
+        # Apply mask (element-wise either input or masking value)
+        masked_data = self.inner_apply_mask(input_data, mask_data)
 
         try:
             write_image(
@@ -67,3 +79,31 @@ def apply_mask(
             )
         except Exception as e:
             raise RuntimeError(f"Error writing output file: {e}") from e
+
+    def inner_apply_mask(
+        self, input_data: np.ndarray, mask_data: np.ndarray
+    ) -> np.ndarray:
+        """
+        Actual application of brain mask implemented as a helper function for consistency
+        between brain extractor and Synthstrip
+
+        Args:
+            input_data (str or Path): image array.
+            mask_path (str or Path): mask array.
+        """
+        # Check that the input and mask have the same shape
+        if input_data.shape != mask_data.shape:
+            raise ValueError("Input image and mask must have the same dimensions.")
+
+        # check whether a global masking value was passed, otherwise choose minimum
+        if self.masking_value is None:
+            current_masking_value = np.min(input_data)
+        else:
+            current_masking_value = (
+                np.array(self.masking_value).astype(input_data.dtype).item()
+            )
+        # Apply mask (element-wise either input or masking value)
+        masked_data = np.where(
+            mask_data.astype(bool), input_data, current_masking_value
+        )
+        return masked_data

brainles_preprocessing/brain_extraction/hd_bet.py

@@ -13,6 +13,15 @@ class Mode(Enum):
 
 
 class HDBetExtractor(BrainExtractor):
+    def __init__(self, masking_value: Optional[Union[int, float]] = None):
+        """
+        Brain extraction HDBet implementation.
+
+        Args:
+            masking_value (Optional[Union[int, float]], optional): global value to be inserted in the masked areas. Default is None which leads to the minimum of each respective image.
+        """
+        super().__init__(masking_value=masking_value)
+
     def extract(
         self,
         input_image_path: Union[str, Path],

brainles_preprocessing/brain_extraction/synthstrip.py

@@ -21,7 +21,9 @@
 
 class SynthStripExtractor(BrainExtractor):
 
-    def __init__(self, border: int = 1):
+    def __init__(
+        self, border: int = 1, masking_value: Optional[Union[int, float]] = None
+    ):
         """
         Brain extraction using SynthStrip with preprocessing conforming to model requirements.
 
@@ -31,9 +33,10 @@ def __init__(self, border: int = 1):
 
         Args:
             border (int): Mask border threshold in mm. Defaults to 1.
-        """
+            masking_value (Optional[Union[int, float]], optional): global value to be inserted in the masked areas. Default is None which leads to the minimum of each respective image.
 
-        super().__init__()
+        """
+        super().__init__(masking_value=masking_value)
         self.border = border
 
     def _setup_model(self, device: torch.device) -> StripModel:
@@ -219,8 +222,8 @@ def extract(
 
         # write the masked output
         img_data = image.get_fdata()
-        bg = np.min([0, img_data.min()])
-        img_data[mask == 0] = bg
+        # actual application of mask through helper function
+        img_data = self.inner_apply_mask(img_data, mask)
         Nifti1Image(img_data, image.affine, image.header).to_filename(
             masked_image_path,
         )