brainles-preprocessing 0.6.0__tar.gz → 0.6.2__tar.gz

{brainles_preprocessing-0.6.0 → brainles_preprocessing-0.6.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: brainles_preprocessing
-Version: 0.6.0
+Version: 0.6.2
 Summary: Tool for preprocessing tasks in biomedical imaging, with a focus on (but not limited to) multi-modal brain MRI
 License: Apache-2.0
 Keywords: brain tumor,glioma,BraTS,skullstripping,brain extraction
@@ -19,13 +19,15 @@ Classifier: Programming Language :: Python :: 3.13
 Provides-Extra: all
 Provides-Extra: itk-elastix
 Provides-Extra: picsl-greedy
+Provides-Extra: synthstrip
 Requires-Dist: antspyx (>=0.5.4,<1.0.0)
 Requires-Dist: auxiliary (>=0.3.1)
 Requires-Dist: brainles_hd_bet (>=0.0.10,<1.0.0)
 Requires-Dist: itk-elastix (>=0.20.0,<0.21.0) ; extra == "all" or extra == "itk-elastix"
+Requires-Dist: loguru (>=0.7.3,<0.8.0)
 Requires-Dist: nibabel (>=3.2.1)
+Requires-Dist: nipreps-synthstrip (>=0.0.1,<0.0.2) ; extra == "all" or extra == "synthstrip"
 Requires-Dist: numpy (>=1.23.0,<2.0.0)
-Requires-Dist: pathlib (>=1.0.1,<2.0.0)
 Requires-Dist: picsl_greedy (>=0.0.6,<0.0.7) ; extra == "all" or extra == "picsl-greedy"
 Requires-Dist: rich (>=13.6.0,<14.0.0)
 Requires-Dist: scipy (<1.14)
@@ -176,17 +178,33 @@ We provide a (WIP) documentation. Have a look [here](https://brainles-preprocess
 Please credit the authors by citing their work.
 
 ### Registration
-We currently provide support for [ANTs](https://github.com/ANTsX/ANTs) (default), [Niftyreg](https://github.com/KCL-BMEIS/niftyreg) (Linux).
+We currently fully support:
+- [ANTs](https://github.com/ANTsX/ANTs) (default)
+- [Niftyreg](https://github.com/KCL-BMEIS/niftyreg) (Linux)
+
+We also offer basic support for:
+- [greedy](https://greedy.readthedocs.io/en/latest/reference.html#greedy-usage) (Optional dependency, install via: `pip install brainles_preprocessing[picsl_greedy]`)
+- [elastix](https://pypi.org/project/itk-elastix/0.13.0/) (Optional dependency, install via: `pip install brainles_preprocessing[itk-elastix]`)
+
+As of now we do not offer inverse transforms for greedy and elastix. Please resort to ANTs or Niftyreg for this.
 
 ### Atlas Reference
-We provide the SRI-24 atlas from this [publication](https://www.ncbi.nlm.nih.gov/pmc/articles/PMC2915788/).
-However, custom atlases in NIfTI format are supported.
+We provide a range of different atlases, namely:
+-  [SRI24](https://www.nitrc.org/frs/download.php/4502/sri24_anatomy_unstripped_nifti.zip) and its [skull-stripped version](https://www.nitrc.org/frs/download.php/4499/sri24_anatomy_nifti.zip)
+-  A slightly modified SRI24 atlas as found in the BraTS challenges (also with a skull-stripped version)
+-  MNI152: [MNI_ICBM_2009c_Nonlinear_Symmetric](https://nist.mni.mcgill.ca/icbm-152-nonlinear-atlases-2009/)
+-  A slightly modified MNI152 atlas as found in the BraTS challenge 
+
+We also support supplying a custom atlas in NIfTI format
+
 
 ### N4 Bias correction
 We currently provide support for N4 Bias correction based on [SimpleITK](https://simpleitk.org/)
 
 ### Brain extraction
-We currently provide support for [HD-BET](https://github.com/MIC-DKFZ/HD-BET).
+We currently support:
+- [HD-BET](https://github.com/MIC-DKFZ/HD-BET)
+- [SynthStrip](https://surfer.nmr.mgh.harvard.edu/docs/synthstrip/) (Optional dependency, install via: `pip install brainles_preprocessing[synthstrip]`)
 
 ### Defacing
 We currently provide support for [Quickshear](https://github.com/nipy/quickshear).

{brainles_preprocessing-0.6.0 → brainles_preprocessing-0.6.2}/README.md

@@ -138,17 +138,33 @@ We provide a (WIP) documentation. Have a look [here](https://brainles-preprocess
 Please credit the authors by citing their work.
 
 ### Registration
-We currently provide support for [ANTs](https://github.com/ANTsX/ANTs) (default), [Niftyreg](https://github.com/KCL-BMEIS/niftyreg) (Linux).
+We currently fully support:
+- [ANTs](https://github.com/ANTsX/ANTs) (default)
+- [Niftyreg](https://github.com/KCL-BMEIS/niftyreg) (Linux)
+
+We also offer basic support for:
+- [greedy](https://greedy.readthedocs.io/en/latest/reference.html#greedy-usage) (Optional dependency, install via: `pip install brainles_preprocessing[picsl_greedy]`)
+- [elastix](https://pypi.org/project/itk-elastix/0.13.0/) (Optional dependency, install via: `pip install brainles_preprocessing[itk-elastix]`)
+
+As of now we do not offer inverse transforms for greedy and elastix. Please resort to ANTs or Niftyreg for this.
 
 ### Atlas Reference
-We provide the SRI-24 atlas from this [publication](https://www.ncbi.nlm.nih.gov/pmc/articles/PMC2915788/).
-However, custom atlases in NIfTI format are supported.
+We provide a range of different atlases, namely:
+-  [SRI24](https://www.nitrc.org/frs/download.php/4502/sri24_anatomy_unstripped_nifti.zip) and its [skull-stripped version](https://www.nitrc.org/frs/download.php/4499/sri24_anatomy_nifti.zip)
+-  A slightly modified SRI24 atlas as found in the BraTS challenges (also with a skull-stripped version)
+-  MNI152: [MNI_ICBM_2009c_Nonlinear_Symmetric](https://nist.mni.mcgill.ca/icbm-152-nonlinear-atlases-2009/)
+-  A slightly modified MNI152 atlas as found in the BraTS challenge 
+
+We also support supplying a custom atlas in NIfTI format
+
 
 ### N4 Bias correction
 We currently provide support for N4 Bias correction based on [SimpleITK](https://simpleitk.org/)
 
 ### Brain extraction
-We currently provide support for [HD-BET](https://github.com/MIC-DKFZ/HD-BET).
+We currently support:
+- [HD-BET](https://github.com/MIC-DKFZ/HD-BET)
+- [SynthStrip](https://surfer.nmr.mgh.harvard.edu/docs/synthstrip/) (Optional dependency, install via: `pip install brainles_preprocessing[synthstrip]`)
 
 ### Defacing
 We currently provide support for [Quickshear](https://github.com/nipy/quickshear).

{brainles_preprocessing-0.6.0 → brainles_preprocessing-0.6.2}/brainles_preprocessing/brain_extraction/brain_extractor.py

@@ -21,8 +21,6 @@ class BrainExtractor:
         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:
         """
@@ -32,7 +30,6 @@ class BrainExtractor:
             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.
         """
@@ -86,11 +83,10 @@ class HDBetExtractor(BrainExtractor):
         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: Union[str, Mode] = Mode.ACCURATE,
         device: Optional[Union[int, str]] = 0,
-        do_tta: Optional[bool] = True,
+        do_tta: bool = True,
+        **kwargs,
     ) -> None:
         # GPU + accurate + TTA
         """
@@ -100,7 +96,6 @@ class HDBetExtractor(BrainExtractor):
             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.

brainles_preprocessing-0.6.2/brainles_preprocessing/brain_extraction/synthstrip.py

@@ -0,0 +1,231 @@
+# Modified from:
+# https://github.com/nipreps/synthstrip/blob/main/nipreps/synthstrip/cli.py
+# Original copyright (c) 2024, NiPreps developers
+# Licensed under the Apache License, Version 2.0
+# Changes made by the BrainLesion Preprocessing team (2025)
+
+from pathlib import Path
+from typing import Optional, Union, cast
+
+import nibabel as nib
+import numpy as np
+import scipy
+import torch
+from nibabel.nifti1 import Nifti1Image
+from nipreps.synthstrip.model import StripModel
+from nitransforms.linear import Affine
+
+from brainles_preprocessing.brain_extraction.brain_extractor import BrainExtractor
+from brainles_preprocessing.utils.zenodo import fetch_synthstrip
+
+
+class SynthStripExtractor(BrainExtractor):
+
+    def __init__(self, border: int = 1):
+        """
+        Brain extraction using SynthStrip with preprocessing conforming to model requirements.
+
+        This is an optional dependency - to use this extractor, you need to install the `brainles_preprocessing` package with the `synthstrip` extra: `pip install brainles_preprocessing[synthstrip]`
+
+        Adapted from https://github.com/nipreps/synthstrip
+
+        Args:
+            border (int): Mask border threshold in mm. Defaults to 1.
+        """
+
+        super().__init__()
+        self.border = border
+
+    def _setup_model(self, device: torch.device) -> StripModel:
+        """
+        Load SynthStrip model and prepare it for inference on the specified device.
+
+        Args:
+            device: Device to load the model onto.
+
+        Returns:
+            A configured and ready-to-use StripModel.
+        """
+        # necessary for speed gains (according to original nipreps authors)
+        torch.backends.cudnn.benchmark = True
+        torch.backends.cudnn.deterministic = True
+
+        with torch.no_grad():
+            model = StripModel()
+            model.to(device)
+            model.eval()
+
+        # Load the model weights
+        weights_folder = fetch_synthstrip()
+        weights = weights_folder / "synthstrip.1.pt"
+        checkpoint = torch.load(weights, map_location=device)
+        model.load_state_dict(checkpoint["model_state_dict"])
+
+        return model
+
+    def _conform(self, input_nii: Nifti1Image) -> Nifti1Image:
+        """
+        Resample the input image to match SynthStrip's expected input space.
+
+        Args:
+            input_nii (Nifti1Image): Input NIfTI image to conform.
+
+        Raises:
+            ValueError: If the input NIfTI image does not have a valid affine.
+
+        Returns:
+            A new NIfTI image with conformed shape and affine.
+        """
+
+        shape = np.array(input_nii.shape[:3])
+        affine = input_nii.affine
+
+        if affine is None:
+            raise ValueError("Input NIfTI image must have a valid affine.")
+
+        # Get corner voxel centers in index coords
+        corner_centers_ijk = (
+            np.array(
+                [
+                    (i, j, k)
+                    for k in (0, shape[2] - 1)
+                    for j in (0, shape[1] - 1)
+                    for i in (0, shape[0] - 1)
+                ]
+            )
+            + 0.5
+        )
+
+        # Get corner voxel centers in mm
+        corners_xyz = (
+            affine
+            @ np.hstack((corner_centers_ijk, np.ones((len(corner_centers_ijk), 1)))).T
+        )
+
+        # Target affine is 1mm voxels in LIA orientation
+        target_affine = np.diag([-1.0, 1.0, -1.0, 1.0])[:, (0, 2, 1, 3)]
+
+        # Target shape
+        extent = corners_xyz.min(1)[:3], corners_xyz.max(1)[:3]
+        target_shape = ((extent[1] - extent[0]) / 1.0 + 0.999).astype(int)
+
+        # SynthStrip likes dimensions be multiple of 64 (192, 256, or 320)
+        target_shape = np.clip(
+            np.ceil(np.array(target_shape) / 64).astype(int) * 64, 192, 320
+        )
+
+        # Ensure shape ordering is LIA too
+        target_shape[2], target_shape[1] = target_shape[1:3]
+
+        # Coordinates of center voxel do not change
+        input_c = affine @ np.hstack((0.5 * (shape - 1), 1.0))
+        target_c = target_affine @ np.hstack((0.5 * (target_shape - 1), 1.0))
+
+        # Rebase the origin of the new, plumb affine
+        target_affine[:3, 3] -= target_c[:3] - input_c[:3]
+
+        nii = Affine(
+            reference=Nifti1Image(
+                np.zeros(target_shape),
+                target_affine,
+                None,
+            ),
+        ).apply(input_nii)
+        return cast(Nifti1Image, nii)
+
+    def _resample_like(
+        self,
+        image: Nifti1Image,
+        target: Nifti1Image,
+        output_dtype: Optional[np.dtype] = None,
+        cval: Union[int, float] = 0,
+    ) -> Nifti1Image:
+        """
+        Resample the input image to match the target's grid using an identity transform.
+
+        Args:
+            image: The image to be resampled.
+            target: The reference image.
+            output_dtype: Output data type.
+            cval: Value to use for constant padding.
+
+        Returns:
+            A resampled NIfTI image.
+        """
+        result = Affine(reference=target).apply(
+            image,
+            output_dtype=output_dtype,
+            cval=cval,
+        )
+        return cast(Nifti1Image, result)
+
+    def extract(
+        self,
+        input_image_path: Union[str, Path],
+        masked_image_path: Union[str, Path],
+        brain_mask_path: Union[str, Path],
+        device: Union[torch.device, str] = "cuda",
+        num_threads: int = 1,
+        **kwargs,
+    ) -> None:
+        """
+        Extract the brain from an input image using SynthStrip.
+
+        Args:
+            input_image_path (Union[str, Path]): Path to the input image.
+            masked_image_path (Union[str, Path]): Path to the output masked image.
+            brain_mask_path (Union[str, Path]): Path to the output brain mask.
+            device (Union[torch.device, str], optional): Device to use for computation. Defaults to "cuda".
+            num_threads (int, optional): Number of threads to use for computation in CPU mode. Defaults to 1.
+
+        Returns:
+            None: The function saves the masked image and brain mask to the specified paths.
+        """
+
+        device = torch.device(device) if isinstance(device, str) else device
+        model = self._setup_model(device=device)
+
+        if device.type == "cpu" and num_threads > 0:
+            torch.set_num_threads(num_threads)
+
+        # normalize intensities
+        image = nib.load(input_image_path)
+        image = cast(Nifti1Image, image)
+        conformed = self._conform(image)
+        in_data = conformed.get_fdata(dtype="float32")
+        in_data -= in_data.min()
+        in_data = np.clip(in_data / np.percentile(in_data, 99), 0, 1)
+        in_data = in_data[np.newaxis, np.newaxis]
+
+        # predict the surface distance transform
+        input_tensor = torch.from_numpy(in_data).to(device)
+        with torch.no_grad():
+            sdt = model(input_tensor).cpu().numpy().squeeze()
+
+        # unconform the sdt and extract mask
+        sdt_target = self._resample_like(
+            Nifti1Image(sdt, conformed.affine, None),
+            image,
+            output_dtype=np.dtype("int16"),
+            cval=100,
+        )
+        sdt_data = np.asanyarray(sdt_target.dataobj).astype("int16")
+
+        # find largest CC (just do this to be safe for now)
+        components = scipy.ndimage.label(sdt_data.squeeze() < self.border)[0]
+        bincount = np.bincount(components.flatten())[1:]
+        mask = components == (np.argmax(bincount) + 1)
+        mask = scipy.ndimage.morphology.binary_fill_holes(mask)
+
+        # write the masked output
+        img_data = image.get_fdata()
+        bg = np.min([0, img_data.min()])
+        img_data[mask == 0] = bg
+        Nifti1Image(img_data, image.affine, image.header).to_filename(
+            masked_image_path,
+        )
+
+        # write the brain mask
+        hdr = image.header.copy()
+        hdr.set_data_dtype("uint8")
+        Nifti1Image(mask, image.affine, hdr).to_filename(brain_mask_path)

{brainles_preprocessing-0.6.0 → brainles_preprocessing-0.6.2}/brainles_preprocessing/modality.py

@@ -4,7 +4,9 @@ import warnings
 from pathlib import Path
 from typing import Dict, Optional, Union
 
+import torch
 from auxiliary.io import read_image, write_image
+from loguru import logger
 
 from brainles_preprocessing.brain_extraction.brain_extractor import BrainExtractor
 from brainles_preprocessing.constants import Atlas, PreprocessorSteps
@@ -15,9 +17,7 @@ from brainles_preprocessing.registration import (  # TODO: this will throw warni
     NiftyRegRegistrator,
 )
 from brainles_preprocessing.registration.registrator import Registrator
-from brainles_preprocessing.utils.zenodo import verify_or_download_atlases
-
-logger = logging.getLogger(__name__)
+from brainles_preprocessing.utils.zenodo import fetch_atlases
 
 
 class Modality:
@@ -599,6 +599,7 @@ class CenterModality(Modality):
         self,
         brain_extractor: BrainExtractor,
         bet_dir_path: Union[str, Path],
+        use_gpu: bool = True,
     ) -> Path:
         """
 
@@ -607,6 +608,7 @@ class CenterModality(Modality):
         Args:
             brain_extractor (BrainExtractor): The brain extractor object.
             bet_dir_path (str or Path): Directory to store brain extraction results.
+            use_gpu (bool): Whether to use GPU for brain extraction if available.
 
         Returns:
             Path: Path to the extracted brain mask.
@@ -618,11 +620,16 @@ class CenterModality(Modality):
         bet = bet_dir_path / f"{self.modality_name}_bet.nii.gz"
         mask_path = bet_dir_path / f"{self.modality_name}_brain_mask.nii.gz"
 
+        device = torch.device(
+            "cuda" if use_gpu and torch.cuda.is_available() else "cpu"
+        )
+
         brain_extractor.extract(
             input_image_path=self.current,
             masked_image_path=bet,
             brain_mask_path=mask_path,
             log_file_path=bet_log,
+            device=device,
         )
 
         # always temporarily store bet image for center modality, since e.g. quickshear defacing could require it
@@ -667,7 +674,7 @@ class CenterModality(Modality):
 
                 # resolve atlas image path
                 if isinstance(defacer.atlas_image_path, Atlas):
-                    atlas_folder = verify_or_download_atlases()
+                    atlas_folder = fetch_atlases()
                     atlas_image_path = atlas_folder / defacer.atlas_image_path.value
                 else:
                     atlas_image_path = Path(defacer.atlas_image_path)

{brainles_preprocessing-0.6.0 → brainles_preprocessing-0.6.2}/brainles_preprocessing/preprocessor/atlas_centric_preprocessor.py

@@ -2,6 +2,8 @@ import shutil
 from pathlib import Path
 from typing import List, Optional, Union
 
+from loguru import logger
+
 from brainles_preprocessing.brain_extraction.brain_extractor import BrainExtractor
 from brainles_preprocessing.constants import Atlas, PreprocessorSteps
 from brainles_preprocessing.defacing import Defacer, QuickshearDefacer
@@ -9,11 +11,7 @@ from brainles_preprocessing.modality import CenterModality, Modality
 from brainles_preprocessing.n4_bias_correction import N4BiasCorrector
 from brainles_preprocessing.preprocessor.preprocessor import BasePreprocessor
 from brainles_preprocessing.registration.registrator import Registrator
-from brainles_preprocessing.utils.logging_utils import LoggingManager
-from brainles_preprocessing.utils.zenodo import verify_or_download_atlases
-
-logging_man = LoggingManager(name=__name__)
-logger = logging_man.get_logger()
+from brainles_preprocessing.utils.zenodo import fetch_atlases
 
 
 class AtlasCentricPreprocessor(BasePreprocessor):
@@ -44,7 +42,7 @@ class AtlasCentricPreprocessor(BasePreprocessor):
         atlas_image_path: Union[str, Path, Atlas] = Atlas.BRATS_SRI24,
         n4_bias_corrector: Optional[N4BiasCorrector] = None,
         temp_folder: Optional[Union[str, Path]] = None,
-        use_gpu: Optional[bool] = None,
+        use_gpu: bool = True,
         limit_cuda_visible_devices: Optional[str] = None,
     ):
         super().__init__(
@@ -60,7 +58,7 @@ class AtlasCentricPreprocessor(BasePreprocessor):
         )
 
         if isinstance(atlas_image_path, Atlas):
-            atlas_folder = verify_or_download_atlases()
+            atlas_folder = fetch_atlases()
             self.atlas_image_path = atlas_folder / atlas_image_path.value
         else:
             self.atlas_image_path = Path(atlas_image_path)
@@ -101,97 +99,102 @@ class AtlasCentricPreprocessor(BasePreprocessor):
         Results are saved in the specified directories, allowing for modular and configurable output storage.
         """
 
-        logging_man._set_log_file(log_file)
-        logger.info(f"{' Starting preprocessing ':=^80}")
-        logger.info(f"Logs are saved to {logging_man.log_file_handler.baseFilename}")
-        modality_names = ", ".join(
-            [modality.modality_name for modality in self.moving_modalities]
-        )
-        logger.info(
-            f"Received center modality: {self.center_modality.modality_name} "
-            f"and moving modalities: {modality_names}"
-        )
-
-        # Co-register moving modalities to center modality
-        logger.info(f"{' Starting Coregistration ':-^80}")
-        self.run_coregistration(
-            save_dir_coregistration=save_dir_coregistration,
-        )
-        logger.info(
-            f"Coregistration complete. Output saved to {save_dir_coregistration}"
-        )
-
-        # Register center modality to atlas
-        logger.info(f"{' Starting atlas registration ':-^80}")
-        self.run_atlas_registration(
-            save_dir_atlas_registration=save_dir_atlas_registration,
-        )
-        logger.info(
-            f"Transformations complete. Output saved to {save_dir_atlas_registration}"
-        )
+        logger_id = self._add_log_file_handler(log_file)
+        try:
+            logger.info(f"{' Starting preprocessing ':=^80}")
+            modality_names = ", ".join(
+                [modality.modality_name for modality in self.moving_modalities]
+            )
+            logger.info(
+                f"Received center modality: {self.center_modality.modality_name} "
+                f"and moving modalities: {modality_names}"
+            )
 
-        # Optional: additional correction in atlas space
-        logger.info(f"{' Checking optional atlas correction ':-^80}")
-        self.run_atlas_correction(
-            save_dir_atlas_correction=save_dir_atlas_correction,
-        )
+            # Co-register moving modalities to center modality
+            logger.info(f"{' Starting Coregistration ':-^80}")
+            self.run_coregistration(
+                save_dir_coregistration=save_dir_coregistration,
+            )
+            logger.info(
+                f"Coregistration complete. Output saved to {save_dir_coregistration}"
+            )
 
-        # Optional: N4 bias correction
-        logger.info(f"{' Checking optional N4 bias correction ':-^80}")
-        self.run_n4_bias_correction(
-            save_dir_n4_bias_correction=save_dir_n4_bias_correction,
-        )
+            # Register center modality to atlas
+            logger.info(f"{' Starting atlas registration ':-^80}")
+            self.run_atlas_registration(
+                save_dir_atlas_registration=save_dir_atlas_registration,
+            )
+            logger.info(
+                f"Transformations complete. Output saved to {save_dir_atlas_registration}"
+            )
 
-        # Now we save images that are not skullstripped (current image = atlas registered or atlas registered + corrected)
-        logger.info("Saving non skull-stripped images...")
-        for modality in self.all_modalities:
-            if modality.raw_skull_output_path:
-                modality.save_current_image(
-                    modality.raw_skull_output_path,
-                    normalization=False,
-                )
-            if modality.normalized_skull_output_path:
-                modality.save_current_image(
-                    modality.normalized_skull_output_path,
-                    normalization=True,
-                )
+            # Optional: additional correction in atlas space
+            logger.info(f"{' Checking optional atlas correction ':-^80}")
+            self.run_atlas_correction(
+                save_dir_atlas_correction=save_dir_atlas_correction,
+            )
 
-        # Optional: Brain extraction
-        logger.info(f"{' Checking optional brain extraction ':-^80}")
-        self.run_brain_extraction(
-            save_dir_brain_extraction=save_dir_brain_extraction,
-        )
+            # Optional: N4 bias correction
+            logger.info(f"{' Checking optional N4 bias correction ':-^80}")
+            self.run_n4_bias_correction(
+                save_dir_n4_bias_correction=save_dir_n4_bias_correction,
+            )
 
-        # Defacing
-        logger.info(f"{' Checking optional defacing ':-^80}")
-        self.run_defacing(
-            save_dir_defacing=save_dir_defacing,
-        )
+            # Now we save images that are not skullstripped (current image = atlas registered or atlas registered + corrected)
+            logger.info("Saving non skull-stripped images...")
+            for modality in self.all_modalities:
+                if modality.raw_skull_output_path:
+                    modality.save_current_image(
+                        modality.raw_skull_output_path,
+                        normalization=False,
+                    )
+                if modality.normalized_skull_output_path:
+                    modality.save_current_image(
+                        modality.normalized_skull_output_path,
+                        normalization=True,
+                    )
+
+            # Optional: Brain extraction
+            logger.info(f"{' Checking optional brain extraction ':-^80}")
+            self.run_brain_extraction(
+                save_dir_brain_extraction=save_dir_brain_extraction,
+            )
 
-        # move to separate method
-        if save_dir_transformations:
-            save_dir_transformations = Path(save_dir_transformations)
+            # Defacing
+            logger.info(f"{' Checking optional defacing ':-^80}")
+            self.run_defacing(
+                save_dir_defacing=save_dir_defacing,
+            )
 
-            # Save transformation matrices
-            logger.info(f"Saving transformation matrices to {save_dir_transformations}")
-            for modality in self.all_modalities:
+            # move to separate method
+            if save_dir_transformations:
+                save_dir_transformations = Path(save_dir_transformations)
 
-                modality_transformations_dir = (
-                    save_dir_transformations / modality.modality_name
+                # Save transformation matrices
+                logger.info(
+                    f"Saving transformation matrices to {save_dir_transformations}"
                 )
-                modality_transformations_dir.mkdir(exist_ok=True, parents=True)
-                for step, path in modality.transformation_paths.items():
-                    if path is not None:
-                        shutil.copyfile(
-                            src=str(path.absolute()),
-                            dst=str(
-                                modality_transformations_dir
-                                / f"{step.value}_{path.name}"
-                            ),
-                        )
-
-        # End
-        logger.info(f"{' Preprocessing complete ':=^80}")
+                for modality in self.all_modalities:
+
+                    modality_transformations_dir = (
+                        save_dir_transformations / modality.modality_name
+                    )
+                    modality_transformations_dir.mkdir(exist_ok=True, parents=True)
+                    for step, path in modality.transformation_paths.items():
+                        if path is not None:
+                            shutil.copyfile(
+                                src=str(path.absolute()),
+                                dst=str(
+                                    modality_transformations_dir
+                                    / f"{step.value}_{path.name}"
+                                ),
+                            )
+
+            # End
+            logger.info(f"{' Preprocessing complete ':=^80}")
+        finally:
+            # Remove log file handler if it was added
+            logger.remove(logger_id)
 
     def run_atlas_registration(
         self, save_dir_atlas_registration: Optional[Union[str, Path]] = None