Implement bias field correction

mri_specific_operations/n4_bias_field_correction.py

+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+Module provides functionality for a N4 bias field correction for MRI images.
+
+Created on Wed Jan 22 14:22:25 2025
+
+@author: jsteiglechner
+"""
+
+from pathlib import Path
+from typing import Union
+
+import matplotlib.pyplot as plt
+import numpy as np
+import SimpleITK as sitk
+
+from ..data_flow.nifti_operations import read_nii, save_nii
+
+
+def perform_n4_correction(
+        input_afp: Union[Path, str],
+        out_afp: Union[Path, str],
+        verbose_plot: bool = False,
+):
+    """
+    Perform N4 intensity bias field correction.
+
+    Notes
+    -----
+    After: https://simpleitk.readthedocs.io/en/master/link_N4BiasFieldCorrection_docs.html
+    https://simpleitk.org/doxygen/latest/html/classitk_1_1simple_1_1N4BiasFieldCorrectionImageFilter.html
+    See also at Slicer: https://www.slicer.org/wiki/Modules:N4ITKBiasFieldCorrection-Documentation-3.6
+
+    Parameters
+    ----------
+    input_afp : Union[Path, str]
+        Path to input image.
+    out_afp : Union[Path, str]
+        Path where corrected image should be saved.
+    verbose_plot : bool, optional
+        Flag whether result should be plotted. The default is False.
+
+    Returns
+    -------
+    None.
+
+    """
+    input_image = sitk.ReadImage(str(input_afp), sitk.sitkFloat32)
+    img = input_image
+
+    # mask = sitk.OtsuThreshold(img, 0, 1, 200)
+    input_array = sitk.GetArrayFromImage(input_image)
+    input_array = np.ones(input_array.shape).astype(np.ubyte)
+
+    mask = sitk.GetImageFromArray(input_array)
+    mask.CopyInformation(input_image)
+    if verbose_plot:
+        plt.imshow(sitk.GetArrayFromImage(mask)[100])
+
+    corrector = sitk.N4BiasFieldCorrectionImageFilter()
+
+    corrector.SetUseMaskLabel(False)
+
+    # Modify corrector setting
+    # number_fitting_levels = 1
+    # maximum_number_of_iterations = 50
+    # corrector.SetMaximumNumberOfIterations([maximum_number_of_iterations] *
+    #                                        number_fitting_levels)
+    # corrector.SetConvergenceThreshold(0.0)
+
+    corrected_image = corrector.Execute(img)  # , mask)
+
+    log_bias_field = corrector.GetLogBiasFieldAsImage(input_image)
+
+    bias_field = input_image / sitk.Exp(log_bias_field)
+
+    bias_field_array = sitk.GetArrayFromImage(bias_field)
+    if verbose_plot:
+        plt.imshow(bias_field_array[100])
+        plt.imshow(sitk.GetArrayFromImage(corrected_image)[:, 100])
+
+    sitk.WriteImage(corrected_image, str(out_afp))
+
+    img, affine = read_nii(out_afp)
+    orig, orig_affine = read_nii(input_afp)
+    if not (np.allclose(affine, orig_affine) and img.shape == orig.shape):
+        print("Physical space is not conserved during N4 correction.")
+        print("Save with original affine. Check this out!")
+        save_nii(img, orig_affine, out_afp)