Merge pull request #90 from yasinzaii/main

Code Updates and Cleaning, Bug Fixes

Author: neuronflow

README.md

@@ -8,11 +8,11 @@
 <!-- [![codecov](https://codecov.io/gh/BrainLesion/brainles-preprocessing/graph/badge.svg?token=A7FWUKO9Y4)](https://codecov.io/gh/BrainLesion/brainles-preprocessing) -->
 <!-- [![License](https://img.shields.io/badge/License-Apache%202.0-blue.svg)](https://opensource.org/licenses/Apache-2.0) -->
 
-`BrainLes preprocessing` is a comprehensive tool for preprocessing tasks in biomedical imaging, with a focus on (but not limited to) multi-modal brain MRI. It can be used to build to build modular preprocessing pipelines:
+`BrainLes preprocessing` is a comprehensive tool for preprocessing tasks in biomedical imaging, with a focus on (but not limited to) multi-modal brain MRI. It can be used to build modular preprocessing pipelines:
 
 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 extration tools.
+BrainLes is written `backend-agnostic` meaning it allows to swap the registration and brain extraction tools.
 
 <!-- TODO mention defacing -->
 

brainles_preprocessing/brain_extraction/brain_extractor.py

@@ -1,90 +1,146 @@
 # TODO add typing and docs
-from abc import abstractmethod
+import shutil
+from abc import ABC, abstractmethod
 from pathlib import Path
-from shutil import copyfile
+from typing import Optional, Union
+from enum import Enum
 
 from auxiliary.nifti.io import read_nifti, write_nifti
-from auxiliary.turbopath import name_extractor
 from brainles_hd_bet import run_hd_bet
 
 
+class Mode(Enum):
+    FAST = "fast"
+    ACCURATE = "accurate"
+
+
 class BrainExtractor:
     @abstractmethod
     def extract(
         self,
-        input_image_path: str,
-        masked_image_path: str,
-        brain_mask_path: str,
-        log_file_path: str,
-        # TODO convert mode to enum
-        mode: str,
+        input_image_path: Union[str, Path],
+        masked_image_path: Union[str, Path],
+        brain_mask_path: Union[str, Path],
+        log_file_path: Optional[Union[str, Path]],
+        mode: Union[str, Mode],
+        **kwargs,
     ) -> None:
+        """
+        Abstract method to extract the brain from an input image.
+
+        Args:
+            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.
+            log_file_path (str or Path, Optional): Path to the log file.
+            mode (str or Mode): Extraction mode.
+            **kwargs: Additional keyword arguments.
+        """
         pass
 
     def apply_mask(
         self,
-        input_image_path: Path,
-        mask_path: Path,
-        bet_image_path: Path,
+        input_image_path: Union[str, Path],
+        mask_path: Union[str, Path],
+        bet_image_path: Union[str, Path],
     ) -> None:
         """
         Apply a brain mask to an input image.
 
         Args:
-            input_image_path (str): Path to the input image (NIfTI format).
-            mask_path (str): Path to the brain mask image (NIfTI format).
-            bet_image_path (str): Path to save the resulting masked image (NIfTI format).
+            input_image_path (str or Path): Path to the input image (NIfTI format).
+            mask_path (str or Path): Path to the brain mask image (NIfTI format).
+            bet_image_path (str or Path): Path to save the resulting masked image (NIfTI format).
         """
 
-        # read data
-        input_data = read_nifti(input_image_path)
-        mask_data = read_nifti(mask_path)
+        try:
+            # Read data
+            input_data = read_nifti(str(input_image_path))
+            mask_data = read_nifti(str(mask_path))
+        except FileNotFoundError as e:
+            raise FileNotFoundError(f"File not found: {e.filename}") from e
+        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
+        # Mask and save it
         masked_data = input_data * mask_data
 
-        write_nifti(
-            input_array=masked_data,
-            output_nifti_path=bet_image_path,
-            reference_nifti_path=input_image_path,
-            create_parent_directory=True,
-        )
+        try:
+            write_nifti(
+                input_array=masked_data,
+                output_nifti_path=str(bet_image_path),
+                reference_nifti_path=str(input_image_path),
+                create_parent_directory=True,
+            )
+        except Exception as e:
+            raise RuntimeError(f"Error writing output file: {e}") from e
 
 
 class HDBetExtractor(BrainExtractor):
     def extract(
         self,
-        input_image_path: str,
-        masked_image_path: str,
-        brain_mask_path: str,
-        log_file_path: str = None,
+        input_image_path: Union[str, Path],
+        masked_image_path: Union[str, Path],
+        brain_mask_path: Union[str, Path],
+        log_file_path: Optional[Union[str, Path]] = None,
         # TODO convert mode to enum
-        mode: str = "accurate",
-        device: int | str = 0,
-        do_tta: bool = True,
+        mode: Union[str, Mode] = Mode.ACCURATE,
+        device: Optional[Union[int, str]] = 0,
+        do_tta: Optional[bool] = True,
     ) -> None:
         # GPU + accurate + TTA
-        """skullstrips images with HD-BET generates a skullstripped file and mask"""
+        """
+        Skull-strips images with HD-BET and generates a skull-stripped file and mask.
+
+        Args:
+            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.
+            log_file_path (str or Path, Optional): Path to the log file.
+            mode (str or Mode): Extraction mode ('fast' or 'accurate').
+            device (str or int): Device to use for computation (e.g., 0 for GPU 0, 'cpu' for CPU).
+            do_tta (bool): whether to do test time data augmentation by mirroring along all axes.
+        """
+
+        # Ensure mode is a Mode enum instance
+        if isinstance(mode, str):
+            try:
+                mode_enum = Mode(mode.lower())
+            except ValueError:
+                raise ValueError(f"'{mode}' is not a valid Mode.")
+        elif isinstance(mode, Mode):
+            mode_enum = mode
+        else:
+            raise TypeError("Mode must be a string or a Mode enum instance.")
+
+        # Run HD-BET
         run_hd_bet(
-            mri_fnames=[input_image_path],
-            output_fnames=[masked_image_path],
-            # device=0,
-            # TODO consider postprocessing
-            # postprocess=False,
-            mode=mode,
+            mri_fnames=[str(input_image_path)],
+            output_fnames=[str(masked_image_path)],
+            mode=mode_enum.value,
             device=device,
+            # TODO consider postprocessing
             postprocess=False,
             do_tta=do_tta,
             keep_mask=True,
             overwrite=True,
         )
 
-        hdbet_mask_path = (
-            Path(masked_image_path).parent
-            / f"{name_extractor(masked_image_path)}_mask.nii.gz"
+        # Construct the path to the generated mask
+        masked_image_path = Path(masked_image_path)
+        hdbet_mask_path = masked_image_path.with_name(
+            masked_image_path.name.replace(".nii.gz", "_mask.nii.gz")
         )
+
         if hdbet_mask_path.resolve() != Path(brain_mask_path).resolve():
-            copyfile(
-                src=hdbet_mask_path,
-                dst=brain_mask_path,
-            )
+            try:
+                shutil.copyfile(
+                    src=str(hdbet_mask_path),
+                    dst=str(brain_mask_path),
+                )
+            except Exception as e:
+                raise RuntimeError(f"Error copying mask file: {e}") from e

brainles_preprocessing/defacing/defacer.py

@@ -1,43 +1,75 @@
-from abc import abstractmethod
+from abc import ABC, abstractmethod
 from pathlib import Path
+from typing import Union
 
 from auxiliary.nifti.io import read_nifti, write_nifti
 
 
-class Defacer:
+class Defacer(ABC):
+    """
+    Base class for defacing medical images using brain masks.
+
+    Subclasses should implement the `deface` method to generate a defaced image
+    based on the provided input image and mask.
+    """
+
     @abstractmethod
     def deface(
         self,
-        input_image_path: Path,
-        mask_image_path: Path,
+        input_image_path: Union[str, Path],
+        mask_image_path: Union[str, Path],
     ) -> None:
+        """
+        Generate a defacing mask provided an input image.
+
+        Args:
+            input_image_path (str or Path): Path to the input image (NIfTI format).
+            mask_image_path (str or Path): Path to the output mask image (NIfTI format).
+        """
         pass
 
     def apply_mask(
         self,
-        input_image_path: str,
-        mask_path: str,
-        defaced_image_path: str,
+        input_image_path: Union[str, Path],
+        mask_path: Union[str, Path],
+        defaced_image_path: Union[str, Path],
     ) -> None:
         """
         Apply a brain mask to an input image.
 
         Args:
-            input_image_path (str): Path to the input image (NIfTI format).
-            mask_path (str): Path to the brain mask image (NIfTI format).
-            defaced_image_path (str): Path to save the resulting defaced image (NIfTI format).
+            input_image_path (str or Path): Path to the input image (NIfTI format).
+            mask_path (str or Path): Path to the brain mask image (NIfTI format).
+            defaced_image_path (str or Path): Path to save the resulting defaced image (NIfTI format).
         """
 
-        # read data
-        input_data = read_nifti(input_image_path)
-        mask_data = read_nifti(mask_path)
+        if not input_image_path.is_file():
+            raise FileNotFoundError(
+                f"Input image file does not exist: {input_image_path}"
+            )
+        if not mask_path.is_file():
+            raise FileNotFoundError(f"Mask file does not exist: {mask_path}")
+
+        try:
+            # Read data
+            input_data = read_nifti(str(input_image_path))
+            mask_data = read_nifti(str(mask_path))
+        except Exception as e:
+            raise RuntimeError(
+                f"An error occurred while reading input 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
+        # Apply mask (element-wise multiplication)
         masked_data = input_data * mask_data
 
+        # Save the defaced image
         write_nifti(
             input_array=masked_data,
-            output_nifti_path=defaced_image_path,
-            reference_nifti_path=input_image_path,
+            output_nifti_path=str(defaced_image_path),
+            reference_nifti_path=str(input_image_path),
             create_parent_directory=True,
         )

brainles_preprocessing/defacing/quickshear/quickshear.py

@@ -1,9 +1,10 @@
 from pathlib import Path
+from typing import Union
 
 import nibabel as nib
+
 from brainles_preprocessing.defacing.defacer import Defacer
 from brainles_preprocessing.defacing.quickshear.nipy_quickshear import run_quickshear
-from numpy.typing import NDArray
 from auxiliary.nifti.io import write_nifti
 
 
@@ -37,17 +38,26 @@ def __init__(self, buffer: float = 10.0):
         super().__init__()
         self.buffer = buffer
 
-    def deface(self, mask_image_path: Path, bet_img_path: Path) -> None:
-        """Deface image using Quickshear algorithm
+    def deface(
+        self,
+        input_image_path: Union[str, Path],
+        mask_image_path: Union[str, Path],
+    ) -> None:
+        """
+        Generate a defacing mask using Quickshear algorithm.
+
+        Note:
+        The input image must be a brain-extracted (skull-stripped) image.
 
         Args:
-            bet_img_path (Path): Path to the brain extracted image
+        input_image_path (str or Path): Path to the brain-extracted input image.
+        mask_image_path (str or Path): Path to save the generated mask image.
         """
 
-        bet_img = nib.load(bet_img_path)
+        bet_img = nib.load(str(input_image_path))
         mask = run_quickshear(bet_img=bet_img, buffer=self.buffer)
         write_nifti(
             input_array=mask,
-            output_nifti_path=mask_image_path,
-            reference_nifti_path=bet_img_path,
+            output_nifti_path=str(mask_image_path),
+            reference_nifti_path=str(input_image_path),
         )