Merge pull request #133 from BrainLesion/78-feature-request-n4-bias-correction-support

78 feature request n4 bias correction support

Author: neuronflow

brainles_preprocessing/brain_extraction/brain_extractor.py

@@ -5,7 +5,7 @@
 from typing import Optional, Union
 from enum import Enum
 
-from auxiliary.nifti.io import read_nifti, write_nifti
+from auxiliary.io import read_image, write_image
 from brainles_hd_bet import run_hd_bet
 
 
@@ -55,8 +55,8 @@ def apply_mask(
 
         try:
             # Read data
-            input_data = read_nifti(str(input_image_path))
-            mask_data = read_nifti(str(mask_path))
+            input_data = read_image(str(input_image_path))
+            mask_data = read_image(str(mask_path))
         except FileNotFoundError as e:
             raise FileNotFoundError(f"File not found: {e.filename}") from e
         except Exception as e:
@@ -70,10 +70,10 @@ def apply_mask(
         masked_data = input_data * mask_data
 
         try:
-            write_nifti(
+            write_image(
                 input_array=masked_data,
-                output_nifti_path=str(bet_image_path),
-                reference_nifti_path=str(input_image_path),
+                output_path=str(bet_image_path),
+                reference_path=str(input_image_path),
                 create_parent_directory=True,
             )
         except Exception as e:

brainles_preprocessing/constants.py

@@ -6,8 +6,9 @@ class PreprocessorSteps(IntEnum):
     COREGISTERED = 1
     ATLAS_REGISTERED = 2
     ATLAS_CORRECTED = 3
-    BET = 4
-    DEFACED = 5
+    N4_BIAS_CORRECTED = 4
+    BET = 5
+    DEFACED = 6
 
 
 class Atlas(str, Enum):

brainles_preprocessing/defacing/defacer.py

@@ -2,7 +2,7 @@
 from pathlib import Path
 from typing import Union
 
-from auxiliary.nifti.io import read_nifti, write_nifti
+from auxiliary.io import read_image, write_image
 
 
 class Defacer(ABC):
@@ -52,8 +52,8 @@ def apply_mask(
 
         try:
             # Read data
-            input_data = read_nifti(str(input_image_path))
-            mask_data = read_nifti(str(mask_path))
+            input_data = read_image(str(input_image_path))
+            mask_data = read_image(str(mask_path))
         except Exception as e:
             raise RuntimeError(
                 f"An error occurred while reading input files: {e}"
@@ -67,9 +67,9 @@ def apply_mask(
         masked_data = input_data * mask_data
 
         # Save the defaced image
-        write_nifti(
+        write_image(
             input_array=masked_data,
-            output_nifti_path=str(defaced_image_path),
-            reference_nifti_path=str(input_image_path),
+            output_path=str(defaced_image_path),
+            reference_path=str(input_image_path),
             create_parent_directory=True,
         )

brainles_preprocessing/defacing/quickshear/quickshear.py

@@ -2,10 +2,10 @@
 from typing import Union
 
 import nibabel as nib
+from auxiliary.io import read_image, write_image
 
 from brainles_preprocessing.defacing.defacer import Defacer
 from brainles_preprocessing.defacing.quickshear.nipy_quickshear import run_quickshear
-from auxiliary.nifti.io import write_nifti
 
 
 class QuickshearDefacer(Defacer):
@@ -56,8 +56,11 @@ def deface(
 
         bet_img = nib.load(str(input_image_path))
         mask = run_quickshear(bet_img=bet_img, buffer=self.buffer)
-        write_nifti(
+
+        # transpose to match simpletik order
+        mask = mask.transpose(2, 1, 0)
+        write_image(
             input_array=mask,
-            output_nifti_path=str(mask_image_path),
-            reference_nifti_path=str(input_image_path),
+            output_path=str(mask_image_path),
+            reference_path=str(input_image_path),
         )

brainles_preprocessing/modality.py

@@ -4,7 +4,7 @@
 from pathlib import Path
 from typing import Dict, Optional, Union
 
-from auxiliary.nifti.io import read_nifti, write_nifti
+from auxiliary.io import read_image, write_image
 
 from brainles_preprocessing.brain_extraction.brain_extractor import BrainExtractor
 from brainles_preprocessing.constants import PreprocessorSteps
@@ -34,6 +34,7 @@ class Modality:
         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.
+        n4_bias_correction (bool, optional): Indicates whether N4 bias correction should be performed.
 
     Attributes:
         modality_name (str): Name of the modality.
@@ -47,6 +48,7 @@ class Modality:
         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.
+        n4_bias_correction (bool): Indicates whether N4 bias correction should be performed.
         coregistration_transform_path (str or None): Path to the coregistration transformation matrix, will be set after coregistration.
 
     Example:
@@ -72,13 +74,15 @@ def __init__(
         normalized_skull_output_path: Optional[Union[str, Path]] = None,
         normalized_defaced_output_path: Optional[Union[str, Path]] = None,
         atlas_correction: bool = True,
+        n4_bias_correction: bool = False,
     ) -> None:
         # Basics
         self.modality_name = modality_name
         self.input_path = Path(input_path)
         self.current = self.input_path
         self.normalizer = normalizer
         self.atlas_correction = atlas_correction
+        self.n4_bias_correction = n4_bias_correction
         self.transformation_paths: Dict[PreprocessorSteps, Path | None] = {}
 
         # Check that atleast one output is generated
@@ -195,12 +199,12 @@ def normalize(
 
         # Normalize the image
         if self.normalizer:
-            image = read_nifti(str(self.current))
+            image = read_image(str(self.current))
             normalized_image = self.normalizer.normalize(image=image)
-            write_nifti(
+            write_image(
                 input_array=normalized_image,
-                output_nifti_path=str(self.current),
-                reference_nifti_path=str(self.current),
+                output_path=str(self.current),
+                reference_path=str(self.current),
             )
         else:
             logger.info("No normalizer specified; skipping normalization.")
@@ -505,12 +509,12 @@ def save_current_image(
         if normalization:
             if self.normalizer is None:
                 raise ValueError("Normalizer is required for normalization.")
-            image = read_nifti(str(self.current))
+            image = read_image(str(self.current))
             normalized_image = self.normalizer.normalize(image=image)
-            write_nifti(
+            write_image(
                 input_array=normalized_image,
-                output_nifti_path=str(output_path),
-                reference_nifti_path=str(self.current),
+                output_path=str(output_path),
+                reference_path=str(self.current),
             )
         else:
             shutil.copyfile(
@@ -534,6 +538,7 @@ class CenterModality(Modality):
         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.
+        n4_bias_correction (bool, optional): Indicates whether N4 bias 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.
 
@@ -549,6 +554,7 @@ class CenterModality(Modality):
         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.
+        n4_bias_correction (bool): Indicates whether N4 bias 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.
 
@@ -575,6 +581,7 @@ def __init__(
         normalized_skull_output_path: Optional[Union[str, Path]] = None,
         normalized_defaced_output_path: Optional[Union[str, Path]] = None,
         atlas_correction: bool = True,
+        n4_bias_correction: bool = False,
         bet_mask_output_path: Optional[Union[str, Path]] = None,
         defacing_mask_output_path: Optional[Union[str, Path]] = None,
     ) -> None:
@@ -589,6 +596,7 @@ def __init__(
             normalized_skull_output_path=normalized_skull_output_path,
             normalized_defaced_output_path=normalized_defaced_output_path,
             atlas_correction=atlas_correction,
+            n4_bias_correction=n4_bias_correction,
         )
         # Only for CenterModality
         self.bet_mask_output_path = (

brainles_preprocessing/n4_bias_correction/__init__.py

@@ -0,0 +1,2 @@
+from .n4_bias_corrector import N4BiasCorrector
+from .sitk.sitk_n4_bias_corrector import SitkN4BiasCorrector

brainles_preprocessing/n4_bias_correction/n4_bias_corrector.py

@@ -0,0 +1,13 @@
+from abc import ABC, abstractmethod
+from typing import Any
+
+
+class N4BiasCorrector(ABC):
+
+    @abstractmethod
+    def correct(
+        self,
+        input_img_path: Any,
+        output_img_path: Any,
+    ) -> None:
+        pass

brainles_preprocessing/n4_bias_correction/sitk/__init__.py

@@ -0,0 +1,82 @@
+from pathlib import Path
+from typing import Callable, Optional, Union
+
+import SimpleITK as sitk
+from auxiliary.io import write_image
+
+from brainles_preprocessing.n4_bias_correction.n4_bias_corrector import N4BiasCorrector
+
+
+class SitkN4BiasCorrector(N4BiasCorrector):
+
+    def __init__(
+        self,
+        mask_func: Optional[Callable[[sitk.Image], sitk.Image]] = None,
+        n_max_iterations: Optional[int] = None,
+        n_fitting_levels: int = 3,
+    ) -> None:
+        """
+        N4 Bias Corrector using SimpleITK.
+
+        Args:
+            mask_func (Optional[Callable[[sitk.Image], sitk.Image]], optional):
+                Function that generates a mask from an image.
+                Defaults to: `lambda img_itk: sitk.OtsuThreshold(img_itk, 0, 1, 200)`.
+            n_max_iterations (Optional[int], optional):
+                Maximum number of iterations for bias field correction.
+            n_fitting_levels (int, optional):
+                Number of fitting levels. Default is 3.
+        """
+
+        if mask_func is None:
+            mask_func = lambda img_itk: sitk.OtsuThreshold(img_itk, 0, 1, 200)
+        self.mask_func = mask_func
+        self.n_max_iterations = n_max_iterations
+        self.n_fitting_levels = n_fitting_levels
+
+    def compute_mask(self, img_itk: sitk.Image) -> sitk.Image:
+        """
+        Compute the mask for the input image using the provided mask function.
+
+        Args:
+            img_itk (SimpleITK.Image): The input image in SimpleITK format.
+
+        Returns:
+            SimpleITK.Image: The computed mask.
+        """
+
+        return self.mask_func(img_itk)
+
+    def correct(
+        self,
+        input_img_path: Union[str, Path],
+        output_img_path: Union[str, Path],
+    ) -> None:
+        """
+        Correct the bias field of the input image using SimpleITK.
+
+        Args:
+            input_img_path (Union[str, Path]): Path to the input image.
+            output_img_path (Union[str, Path]): Path where the corrected image will be saved.
+
+        Returns:
+            None
+        """
+        img_itk = sitk.ReadImage(str(input_img_path))
+
+        mask_itk = self.compute_mask(img_itk)
+
+        corrector = sitk.N4BiasFieldCorrectionImageFilter()
+        if self.n_max_iterations is not None:
+            corrector.SetMaximumNumberOfIterations(
+                [self.n_max_iterations] * self.n_fitting_levels
+            )
+
+        corrected_img = corrector.Execute(img_itk, mask_itk)
+        corrected_img = sitk.GetArrayFromImage(corrected_img)
+
+        write_image(
+            input_array=corrected_img,
+            output_path=str(output_img_path),
+            reference_path=str(input_img_path),
+        )