diff --git a/brainles_preprocessing/brain_extraction/brain_extractor.py b/brainles_preprocessing/brain_extraction/brain_extractor.py index 908982a..5efa1c4 100644 --- a/brainles_preprocessing/brain_extraction/brain_extractor.py +++ b/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 diff --git a/brainles_preprocessing/brain_extraction/hd_bet.py b/brainles_preprocessing/brain_extraction/hd_bet.py index cd631ea..7391efb 100644 --- a/brainles_preprocessing/brain_extraction/hd_bet.py +++ b/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], diff --git a/brainles_preprocessing/brain_extraction/synthstrip.py b/brainles_preprocessing/brain_extraction/synthstrip.py index 0ea3135..58c8198 100644 --- a/brainles_preprocessing/brain_extraction/synthstrip.py +++ b/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, )