MNutils 1.2.0__py3-none-any.whl

mnutils/utils/images.py

@@ -0,0 +1,202 @@
+import os
+from pathlib import Path
+from typing import TYPE_CHECKING, cast, overload
+
+# This block is ignored at runtime. I only want to import later to reduce overhead.
+if TYPE_CHECKING:
+    import HD_BET.checkpoint_download
+    import HD_BET.hd_bet_prediction
+    import torch
+
+from loguru import logger
+from pandas.compat._optional import import_optional_dependency
+
+from . import file_helpers
+
+
+@overload
+def extract_brain(
+    file: str | Path,
+    device: str = "cuda",
+    use_tta: bool = True,
+    save_mask: bool = True,
+    extract_brain: bool = True,
+    verbose: bool = False,
+) -> Path: ...
+
+
+@overload
+def extract_brain(
+    file: list[str | Path],
+    device: str = "cuda",
+    use_tta: bool = True,
+    save_mask: bool = True,
+    extract_brain: bool = True,
+    verbose: bool = False,
+) -> list[Path]: ...
+
+
+def extract_brain(
+    file: str | Path | list[str | Path],
+    device: str = "cuda",
+    use_tta: bool = True,
+    save_mask: bool = True,
+    extract_brain: bool = True,
+    verbose: bool = False,
+) -> Path | list[Path]:
+    """Extract the brain from an image file using HD-BET.
+
+    This function processes NIfTI files or directories containing DICOM files to extract the
+    brain region using the HD-BET (HD-Brain Extraction Tool). It supports optional test-time
+    augmentation (TTA) and allows saving the brain mask and/or the brain-extracted image.
+
+    Parameters
+    ----------
+    file : str | Path | list[str | Path]
+        The path(s) to the image file(s) or directories containing DICOM files to extract the
+        brain from. If a directory is provided, it will be converted to a NIfTI file first.
+    device : str, optional
+        The device to use for processing. Default is "cuda".
+    use_tta : bool, optional
+        Whether to use test-time augmentation (TTA) during brain extraction. Default is True.
+    save_mask : bool, optional
+        Whether to save the brain mask. Default is True.
+    extract_brain : bool, optional
+        Whether to compute and save the brain-extracted image. Default is True.
+    verbose : bool, optional
+        Whether to enable verbose logging. Default is False.
+
+    Returns
+    -------
+    Path | list[Path]
+        The path(s) to the brain-extracted NIfTI file(s). If a single file is provided as input,
+        a single Path object is returned. If a list of files is provided, a list of Path objects
+        is returned.
+
+    Raises
+    ------
+    ValueError
+        If the input file does not have a ".nii.gz" suffix or if the input is invalid.
+    ImportError
+        If the required dependencies (torch, HD-BET) are not installed.
+
+    Notes
+    -----
+    - This function requires the HD-BET library and PyTorch to be installed.
+    - If the input is a directory containing DICOM files, it will be converted to a NIfTI file
+      before processing.
+    - The function checks for existing brain-extracted files in the output directory to avoid
+      overwriting existing results.
+    - The HD-BET model parameters will be downloaded automatically if not already available.
+
+    Examples
+    --------
+    Extract the brain from a single NIfTI file:
+    >>> extract_brain("path/to/image.nii.gz")
+    Extract the brain from multiple NIfTI files:
+    >>> extract_brain(["path/to/image1.nii.gz", "path/to/image2.nii.gz"])
+    Extract the brain from a directory containing DICOM files:
+    >>> extract_brain("path/to/dicom_directory")
+
+    """
+    # If the the input is a folder containing dicoms first convert to nifti
+    if isinstance(file, (str, Path)):
+        single_file = True
+        file = [Path(file)]
+    else:
+        single_file = False
+
+    files: list[Path] = [Path(f) for f in file]
+
+    for idx, input_path in enumerate(files):
+        if os.path.isdir(input_path):
+            nifti_path = file_helpers.get_nifti_data(input_path, convert_dicoms=True)
+            # Replace the original path with the new nifti path
+            files[idx] = nifti_path
+
+    # type hinting
+    torch = import_optional_dependency("torch")
+    torch = cast("torch", torch)
+    hd_bet_checkpoint_download = import_optional_dependency("HD_BET.checkpoint_download")
+    hd_bet_checkpoint_download = cast("HD_BET.checkpoint_download", hd_bet_checkpoint_download)
+    hd_bet_predictor = import_optional_dependency("HD_BET.hd_bet_prediction")
+    hd_bet_predictor = cast("HD_BET.hd_bet_prediction", hd_bet_predictor)
+
+    # Start running hd_bet
+    # TODO fix issues with stdout capture
+
+    downloaded_parameters = False
+    predictor = None
+
+    output_files: list[Path] = []
+
+    for idx, input_path in enumerate(files):
+        logger.debug(
+            f"Processing file {file_helpers._relative_to_cwd(input_path)} ({idx + 1}/{len(files)})"
+        )
+        # Assert that input path has nii.gz suffix
+        if not str(input_path).endswith(".nii.gz"):
+            raise ValueError(
+                f"Input file {input_path} does not have a .nii.gz suffix. "
+                "Please provide a valid NIfTI file."
+            )
+
+        output_dir = Path(input_path).parent
+        default_output_path = (
+            output_dir / f"{str(input_path.name)[:-7]}_bet.nii.gz"
+        )  # Bet will be saved with _bet suffix by default
+        other_bet_niftis = file_helpers.get_nifti_file(
+            input_path.parent,
+            filter_out_keywords=[],
+            filter_in_keywords=["_bet.nii.gz"],
+        )
+
+        # Check if output file already exists
+        if default_output_path.exists():
+            logger.warning(
+                f"Output file {file_helpers._relative_to_cwd(default_output_path)} "
+                f"already exists. Skipping brain extraction for "
+                f"{file_helpers._relative_to_cwd(input_path)}."
+            )
+            output_files.append(default_output_path)
+            continue
+
+        if other_bet_niftis:
+            logger.warning(
+                f"Other brain extracted NIfTI files found in output directory "
+                f"{file_helpers._relative_to_cwd(output_dir)}. This may indicate that "
+                f"brain extraction has already been performed for this or other files. Skipping "
+                f"brain extraction for {file_helpers._relative_to_cwd(input_path)} to "
+                f"avoid overwriting existing files."
+            )
+            output_files.append(other_bet_niftis)
+            continue
+
+        if not downloaded_parameters:
+            hd_bet_checkpoint_download.maybe_download_parameters()
+            downloaded_parameters = True
+
+        if predictor is None:
+            predictor = hd_bet_predictor.get_hdbet_predictor(
+                use_tta=use_tta, device=torch.device(device), verbose=verbose
+            )
+
+        hd_bet_predictor.hdbet_predict(
+            str(input_path),
+            str(input_path),
+            predictor,
+            keep_brain_mask=save_mask,
+            compute_brain_extracted_image=extract_brain,
+        )
+
+        logger.debug(
+            f"Brain extracted successfully for "
+            f"{file_helpers._relative_to_cwd(input_path)}, saved to "
+            f"{file_helpers._relative_to_cwd(default_output_path)}"
+        )
+        output_files.append(default_output_path)
+
+    if single_file:
+        return output_files[0]
+    else:
+        return output_files

mnutils/utils/nifti.py

@@ -0,0 +1,379 @@
+from __future__ import annotations
+
+from copy import deepcopy
+from typing import Any, Literal
+
+import nibabel as nib
+import numpy as np
+import numpy.typing as npt
+from loguru import logger
+from nibabel import affines, processing, spatialimages
+
+from .. import GESeries
+
+DEFAULT_PARAMS = {
+    "orientation": ("L", "P", "S"),
+}
+
+DEFAULT_RESAMPLE_PARAMS = {
+    "order": 0,
+    "mode": "grid-constant",
+    "cval": np.nan,
+}
+
+type DISPLAY_PLANES = Literal["axial", "coronal", "sagittal"]
+
+
+def get_display_affine(
+    nii: spatialimages.SpatialImage | GESeries.NiiBase,
+    orientation: tuple[str, str, str] = ("L", "P", "S"),
+) -> npt.NDArray[np.float64]:
+    """Return the affine mapping display voxel indices (after orient_nifti) to world mm.
+
+    orient_nifti applies apply_orientation then flipud(rot90(k=1)) which equals
+    a transpose of axes 0 and 1. This function encodes both steps as an affine.
+    Use inv(get_display_affine(t1)) @ mrsi_affine to get mrsi→display mapping.
+    """
+    if isinstance(nii, GESeries.NiiBase):
+        nii: spatialimages.SpatialImage = nii.nii
+    current_ornt = nib.orientations.io_orientation(nii.affine)
+    target_ornt = nib.orientations.axcodes2ornt(orientation)
+    transform = nib.orientations.ornt_transform(current_ornt, target_ornt)
+    reoriented = nii.as_reoriented(transform)  # pyright: ignore[reportArgumentType]
+    oriented_affine = np.asarray(reoriented.affine, dtype=float)
+    # flipud(rot90(k=1, axes=(0,1))) = pure transpose of axes 0 and 1.
+    # This swaps columns 0 and 1 of the affine.
+    P = np.eye(4)
+    P[0, 0] = P[1, 1] = 0.0
+    P[0, 1] = P[1, 0] = 1.0
+    return oriented_affine @ P
+
+
+def orient_nifti(
+    nii: spatialimages.SpatialImage | GESeries.NiiBase,
+    orientation: tuple[str, str, str] | None = None,
+    display_plane: DISPLAY_PLANES | None = None,
+    caching: bool = True,
+) -> npt.NDArray[np.float64]:
+    """Reorient and rotate a NIfTI image's data array for display.
+
+    Applies the requested axis orientation (or one derived from
+    `display_plane`), then transposes axes 0 and 1 to match the display
+    convention used elsewhere in this module.
+
+    Parameters
+    ----------
+    nii : spatialimages.SpatialImage or GESeries.NiiBase
+        The image to orient.
+    orientation : tuple[str, str, str], optional
+        Target axis codes. Defaults to `DEFAULT_PARAMS["orientation"]` if
+        neither this nor `display_plane` is given.
+    display_plane : {"axial", "coronal", "sagittal"}, optional
+        If given, overrides `orientation` with a preset for the named plane.
+    caching : bool, optional
+        Whether to cache the underlying image data (more memory, faster
+        repeat access) or not (less memory, slower repeat access).
+        Defaults to True.
+
+    Returns
+    -------
+    ndarray
+        The reoriented image data array.
+    """
+    if isinstance(nii, GESeries.NiiBase):
+        nii: spatialimages.SpatialImage = nii.nii
+    if display_plane is not None:
+        match display_plane:
+            case "axial":
+                orientation = ("L", "P", "S")
+            case "coronal":
+                orientation = ("L", "I", "P")
+            case "sagittal":
+                orientation = ("I", "L", "P")
+            case _:
+                orientation = DEFAULT_PARAMS["orientation"]
+        logger.debug(f"Orienting NIfTI to {display_plane} plane. Using orientation {orientation}.")
+
+    if orientation is None:
+        logger.debug(
+            f"No orientation specified. Using default orientation: {DEFAULT_PARAMS['orientation']}."
+        )
+        orientation = DEFAULT_PARAMS["orientation"]
+
+    if caching:
+        logger.debug(
+            "Using caching for NIfTI data. This may use more memory but can speed up "
+            "processing for repeat access."
+        )
+        caching_param = "fill"
+    else:
+        logger.debug(
+            "Not using caching for NIfTI data. This may use less memory but can slow down "
+            "processing for repeat access."
+        )
+        caching_param = "unchanged"
+
+    imagesOr = nib.orientations.apply_orientation(
+        nii.get_fdata(caching=caching_param),
+        nib.orientations.axcodes2ornt(orientation),
+    )
+    images = np.flipud(np.rot90(imagesOr, k=1, axes=(0, 1)))
+    return images
+
+
+def apply_half_voxel_shift(
+    affine: npt.NDArray[np.float64], shift_z_axis: bool = False
+) -> npt.NDArray[np.float64]:
+    """Shift an affine by half a voxel in x/y (and optionally z). Deprecated, always raises.
+
+    Parameters
+    ----------
+    affine : ndarray
+        The affine to shift.
+    shift_z_axis : bool, optional
+        Whether to also shift the z axis. Defaults to False.
+
+    Returns
+    -------
+    ndarray
+        The shifted affine. Unreachable, since this function always raises.
+
+    Raises
+    ------
+    DeprecationWarning
+        Always. Half voxel shifts are usually a symptom of a bad affine
+        definition; use dcm2nii to convert DICOMs and use its affine instead.
+    """
+    raise DeprecationWarning(
+        "This function should not be used as half voxel shifts are usually bad and a "
+        "result of bad affine definition"
+        "Use dcm2nii to convert dicoms to nifits and use the resulting affine"
+    )
+    # Half voxel shift should only be in x and y directions
+    if shift_z_axis:
+        dims = 3
+    else:
+        dims = 2
+
+    shifted_affine = deepcopy(affine)
+    shift = np.diag(affine[:dims, :dims]) * 0.5
+    shifted_affine[:dims, 3] += shift
+    logger.debug(f"Applied half voxel shift of {shift} to affine.")
+    return shifted_affine
+
+
+def resample_nifti(
+    source_nii: spatialimages.SpatialImage | GESeries.NiiBase,
+    target_nii: spatialimages.SpatialImage | GESeries.NiiBase,
+    **kwargs: Any,
+) -> spatialimages.SpatialImage:
+    """Resample `source_nii` into `target_nii`'s voxel grid via `nibabel.processing`.
+
+    Parameters
+    ----------
+    source_nii : spatialimages.SpatialImage or GESeries.NiiBase
+        The image to resample.
+    target_nii : spatialimages.SpatialImage or GESeries.NiiBase
+        The image whose grid to resample onto.
+    **kwargs
+        Overrides for `DEFAULT_RESAMPLE_PARAMS`, passed to
+        `nibabel.processing.resample_from_to`.
+
+    Returns
+    -------
+    spatialimages.SpatialImage
+        The resampled image.
+    """
+    if isinstance(source_nii, GESeries.NiiBase):
+        source_nii: spatialimages.SpatialImage = source_nii.nii
+    if isinstance(target_nii, GESeries.NiiBase):
+        target_nii: spatialimages.SpatialImage = target_nii.nii
+    # Resample data_nifti to base_nifti space
+    resample_kwargs = DEFAULT_RESAMPLE_PARAMS.copy()
+    resample_kwargs.update(kwargs)
+    logger.debug(f"Resampling with parameters: {resample_kwargs}")
+
+    resampled_nii: spatialimages.SpatialImage = processing.resample_from_to(
+        source_nii, target_nii, **resample_kwargs
+    )
+    return resampled_nii
+
+
+def resample_and_orient_nifti(
+    source_nii: spatialimages.SpatialImage | GESeries.NiiBase,
+    target_nii: spatialimages.SpatialImage | GESeries.NiiBase,
+    orientation: tuple[str, str, str] | None = DEFAULT_PARAMS["orientation"],
+    display_plane: DISPLAY_PLANES | None = None,
+    **kwargs: Any,
+) -> tuple[npt.NDArray[np.float64], npt.NDArray[np.float64]]:
+    """Resample `source_nii` onto `target_nii`'s grid, then orient both for display.
+
+    Parameters
+    ----------
+    source_nii : spatialimages.SpatialImage or GESeries.NiiBase
+        The image to resample.
+    target_nii : spatialimages.SpatialImage or GESeries.NiiBase
+        The image whose grid to resample onto.
+    orientation : tuple[str, str, str], optional
+        Target axis codes passed to `orient_nifti`. Defaults to
+        `DEFAULT_PARAMS["orientation"]`.
+    display_plane : {"axial", "coronal", "sagittal"}, optional
+        If given, overrides `orientation` with a preset for the named plane.
+    **kwargs
+        Additional keyword arguments passed to `resample_nifti`.
+
+    Returns
+    -------
+    tuple[ndarray, ndarray]
+        The oriented target image data and the oriented resampled data.
+    """
+    if isinstance(source_nii, GESeries.NiiBase):
+        source_nii: spatialimages.SpatialImage = source_nii.nii
+    if isinstance(target_nii, GESeries.NiiBase):
+        target_nii: spatialimages.SpatialImage = target_nii.nii
+    # Resample base_nifti to target_nifti space
+    resampled_nii = resample_nifti(source_nii=source_nii, target_nii=target_nii, **kwargs)
+    # Orient the data
+    base_images = orient_nifti(target_nii, orientation, display_plane=display_plane)
+    data_images = orient_nifti(resampled_nii, orientation, display_plane=display_plane)
+    return base_images, data_images
+
+
+def downsample_to_coverage(
+    src_data: npt.NDArray | spatialimages.SpatialImage | GESeries.NiiBase,
+    tgt_data: spatialimages.SpatialImage | GESeries.NiiBase | None = None,
+    src_affine: npt.NDArray | None = None,
+    tgt_affine: npt.NDArray | None = None,
+    tgt_shape: tuple[int, int, int] | None = None,
+    super_sampling: bool = True,
+    sampling_order: int | list[int] | tuple[int, int, int] | None = None,
+    **kwargs,
+) -> npt.NDArray[np.float32]:
+    """Compute, per target voxel, the fraction of sub-samples that fall inside the source image.
+
+    For each target voxel, samples a supersampling grid of sub-voxel points
+    (or a single center point if `super_sampling` is False), maps them into
+    source voxel space via the affines, and averages how many land inside
+    the source image bounds.
+
+    Parameters
+    ----------
+    src_data : ndarray, spatialimages.SpatialImage, or GESeries.NiiBase
+        The source image data or image. If a plain array, `src_affine` must
+        be given.
+    tgt_data : spatialimages.SpatialImage or GESeries.NiiBase, optional
+        The target image, used to derive `tgt_affine`/`tgt_shape` if not
+        given directly.
+    src_affine : ndarray, optional
+        The source image's affine. Required if `src_data` is a plain array.
+    tgt_affine : ndarray, optional
+        The target image's affine. Required if `tgt_data` is not given.
+    tgt_shape : tuple[int, int, int], optional
+        The target image's shape. Required if `tgt_data` is not given.
+    super_sampling : bool, optional
+        Whether to sample a sub-voxel grid per target voxel (more accurate)
+        or a single center point. Defaults to True.
+    sampling_order : int, list[int], or tuple[int, int, int], optional
+        The supersampling grid size per axis. Defaults to a size derived
+        from the target-to-source voxel scale.
+    **kwargs
+        Unused, accepted for forward compatibility.
+
+    Returns
+    -------
+    ndarray
+        The coverage fraction for each target voxel, shape `tgt_shape`.
+
+    Raises
+    ------
+    ValueError
+        If the source/target affines or target shape cannot be determined.
+    """
+    if isinstance(src_data, GESeries.NiiBase):
+        src_img = src_data.images()
+        src_affine = src_data.nii.affine
+    elif isinstance(src_data, spatialimages.SpatialImage):
+        src_img = src_data.get_fdata()
+        src_affine = src_data.affine
+    else:
+        src_img = src_data
+
+    if isinstance(tgt_data, GESeries.NiiBase):
+        tgt_affine = tgt_data.nii.affine
+        tgt_shape = tgt_data.images().shape
+    elif isinstance(tgt_data, spatialimages.SpatialImage):
+        tgt_affine = tgt_data.affine
+        tgt_shape = tgt_data.get_fdata().shape
+
+    if src_affine is None or tgt_affine is None or tgt_shape is None:
+        raise ValueError(
+            "Source and target affines and target shape must be provided. Either pass "
+            "target data as Nifti/GESeries or provide affines and shape directly."
+        )
+
+    # Generally we do not need to do a half voxel shift as our offsets center the samples
+    # within the voxels
+    # Double check by overlaying on the anatomical again
+    tgt2src_affine = np.linalg.inv(src_affine).dot(tgt_affine)
+
+    logger.debug(f"Source shape: {src_img.shape}, Target shape: {tgt_shape}")
+    logger.debug(f"Transformation from target to source affine:\n{tgt2src_affine}")
+
+    # Supersampling
+
+    if super_sampling:
+        if sampling_order is not None:
+            if isinstance(sampling_order, int):
+                s = [sampling_order] * 3
+            else:
+                s = list(sampling_order)
+        else:
+            s = np.ceil(np.diag(tgt2src_affine)[:3]).astype(int).tolist()
+        logger.debug(f"Using supersampling grid of size: {s}")
+
+        offset = np.multiply(s, [0, 0, -1])
+        logger.debug(f"Offset for sampling: {offset}")
+        # I am not entirely sure why we need to offset z direction by the sampling coefficient...
+        offsets = (np.indices(s).transpose(1, 2, 3, 0) + offset) / s
+    else:
+        logger.warning(
+            "No supersampling. This may lead to inaccurate coverage estimates, especially "
+            "for small voxels or large transformations."
+        )
+        offsets = np.array([[[[-0, -0, -0]]]])  # Does not seem to neet a offest..
+        logger.debug(f"No supersampling. Using single sample at offset: {offsets}")
+
+    # MAIN LOOP
+
+    coverage = np.zeros(tgt_shape, dtype=np.float32)
+
+    for i in range(tgt_shape[0]):
+        for j in range(tgt_shape[1]):
+            for k in range(tgt_shape[2]):
+                # --- Compute index of sub-samples within target voxel ---
+                # target index coordinate + offsets
+                tgt_idx = np.array([i, j, k])[None, None, None, :] + offsets
+                src_idx = affines.apply_affine(tgt2src_affine, tgt_idx)
+
+                # Nearest neighbor sample
+                xi = np.rint(src_idx[..., 0]).astype(int)
+                yi = np.rint(src_idx[..., 1]).astype(int)
+                zi = np.rint(src_idx[..., 2]).astype(int)
+
+                # Boundary mask
+                inside = (
+                    (xi >= 0)
+                    & (xi < src_img.shape[0])
+                    & (yi >= 0)
+                    & (yi < src_img.shape[1])
+                    & (zi >= 0)
+                    & (zi < src_img.shape[2])
+                )
+
+                samples = np.zeros_like(xi, dtype=np.float32)
+                samples[inside] = src_img[xi[inside], yi[inside], zi[inside]].astype(np.float32)
+
+                # Average over samples → coverage fraction
+                coverage[i, j, k] = samples.mean()
+
+    return coverage.astype(np.float32)