bsplot 0.0.2__py3-none-any.whl

bsplot/volume.py

@@ -0,0 +1,958 @@
+import json
+import os
+
+import matplotlib.pyplot as plt
+import nibabel as nib
+import numpy as np
+import numpy.ma as ma
+from matplotlib import colors, patches, transforms
+from scipy.ndimage import center_of_mass
+from skimage.measure import find_contours
+
+from bsplot.colors import Colormap
+
+
+def get_center_of_mass(img, mode="absolute", space="voxel", view=None):
+    """
+    Compute center of mass of image data.
+
+    Parameters
+    ----------
+    img : str, PathLike, or Nifti1Image
+        NIfTI image or path to one.
+    mode : str
+        Weighting mode for center of mass calculation:
+        - 'absolute': use absolute values (default)
+        - 'positive': only positive values
+        - 'negative': only negative values (inverted)
+        - 'relative': shift minimum to zero
+    space : str
+        - 'voxel': return [i, j, k] voxel indices (int)
+        - 'world' or 'mm': return [x, y, z] in millimeters (float)
+    view : str, optional
+        If provided, return only the coordinate for this view plane:
+        - 'sagittal'/'lateral': X coordinate (index 0)
+        - 'coronal'/'frontal': Y coordinate (index 1)
+        - 'axial'/'horizontal'/'transverse': Z coordinate (index 2)
+
+    Returns
+    -------
+    coords : ndarray or float
+        Center of mass coordinates. If view is None, returns shape (3,) array.
+        If view is specified, returns single coordinate (int for voxel, float for world).
+    """
+    if isinstance(img, (str, os.PathLike)):
+        img = nib.load(img)
+
+    data = img.get_fdata()
+    data = np.nan_to_num(data, nan=0.0)
+
+    if mode == "absolute":
+        weights = np.abs(data)
+    elif mode == "positive":
+        weights = np.where(data > 0, data, 0)
+    elif mode == "negative":
+        weights = np.where(data < 0, -data, 0)
+    elif mode == "relative":
+        data = data - np.nanmin(data)
+        weights = np.where(data == 0, 0, data)
+    else:
+        raise ValueError(f"Invalid mode: {mode}")
+
+    if not np.any(weights):
+        com_voxel = np.array([0.0, 0.0, 0.0])
+    else:
+        com_voxel = np.array(center_of_mass(weights))
+
+    if space == "voxel":
+        # Return integer voxel indices, clamped to valid range
+        coords = np.array(
+            [
+                min(max(0, int(round(c))), s - 1)
+                for c, s in zip(com_voxel, img.shape[:3])
+            ]
+        )
+    elif space in ("world", "mm"):
+        coords = nib.affines.apply_affine(img.affine, com_voxel)
+    else:
+        raise ValueError(f"Invalid space: {space}. Use 'voxel' or 'world'/'mm'.")
+
+    if view is not None:
+        idx = get_view_index(view)
+        return int(round(coords[idx])) if space == "voxel" else coords[idx]
+
+    return coords
+
+
+# Backwards compatibility aliases
+def get_cut_coords(nii_img):
+    """Deprecated: use get_center_of_mass(img, space='world')"""
+    return get_center_of_mass(nii_img, space="world", view=None)
+
+
+def get_com_slice_index(nii_img, plane="sagittal"):
+    """Deprecated: use get_center_of_mass(img, space='world', view=plane)"""
+    return get_center_of_mass(nii_img, space="world", view=plane)
+
+
+def get_slice_index(img, mode="absolute"):
+    """Deprecated: use get_center_of_mass(img, mode=mode, space='voxel')"""
+    return list(get_center_of_mass(img, mode=mode, space="voxel", view=None))
+
+
+def get_world_coords(img):
+    """Get world coordinate arrays for each axis."""
+    affine = img.affine
+    shape = img.shape
+
+    x_voxel = np.arange(shape[0])
+    y_voxel = np.arange(shape[1])
+    z_voxel = np.arange(shape[2])
+
+    x_world = affine[0, 0] * x_voxel + affine[0, 3]
+    y_world = affine[1, 1] * y_voxel + affine[1, 3]
+    z_world = affine[2, 2] * z_voxel + affine[2, 3]
+
+    return x_world, y_world, z_world
+
+
+def _normalize_view(view):
+    """Map plot_surf-style view names to canonical slice plane names.
+
+    Accepts all view aliases used by plot_surf and maps them to:
+    'sagittal', 'coronal', or 'horizontal'.
+    """
+    v = view.lower().strip()
+    if v in ("sagittal", "lateral", "medial"):
+        return "sagittal"
+    elif v in ("coronal", "frontal", "anterior", "posterior", "front", "back"):
+        return "coronal"
+    elif v in (
+        "horizontal", "axial", "transverse",
+        "dorsal", "ventral", "top", "bottom", "superior", "inferior",
+    ):
+        return "horizontal"
+    else:
+        raise ValueError(
+            f"Invalid view: '{view}'. Use one of: sagittal/lateral/medial, "
+            f"coronal/frontal/anterior/posterior/front/back, "
+            f"horizontal/axial/transverse/dorsal/ventral/top/bottom/superior/inferior."
+        )
+
+
+def get_view_index(view):
+    """Map view name to axis index (0=sagittal/x, 1=coronal/y, 2=axial/z)."""
+    canon = _normalize_view(view)
+    if canon == "sagittal":
+        return 0
+    elif canon == "coronal":
+        return 1
+    else:
+        return 2
+
+
+def get_plane_indices(view):
+    """Map view name to the two in-plane axis indices."""
+    canon = _normalize_view(view)
+    if canon == "sagittal":
+        return (1, 2)
+    elif canon == "coronal":
+        return (0, 2)
+    else:
+        return (0, 1)
+
+
+def plot_slice(
+    img,
+    ax=None,
+    cmap="grey",
+    view="sagittal",
+    alpha=1.0,
+    sl=None,
+    slice_mm=None,
+    center_coordinate=None,
+    threshold=0,
+    plot_cut_coord=False,
+    annotate_slice_mm=False,  # new: annotate the slice position in mm
+    intensity_projection=None,  # max, mean, sum, absmax
+    diverging=False,
+    outline=False,
+    outline_kwargs={"color": "k", "linewidth": 0.5, "alpha": 1},
+    colorbar=False,
+    set_axis_extent=True,
+    rotate=0,  # rotation in degrees (0, 90, 180, 270)
+    flip_lr=False,  # flip data left-right for radiological convention
+    orientation_labels=None,  # None=no labels, "axes"=axis labels, "corner"=corner indicators
+    hemi=None,  # hemisphere hint for view-dependent x-axis direction
+    hide_hemi=None,  # None, 'left', or 'right' – mask out that hemisphere (x<0 or x>0)
+    **kwargs,
+):
+    # Save the original directional view name before normalising so we
+    # can determine whether the x-extent must be negated to match the
+    # world-coordinate convention used by plot_surf.
+    _original_view = view.lower().strip() if isinstance(view, str) else str(view)
+
+    # Normalize view aliases (e.g. "anterior" → "coronal") so all
+    # downstream checks use canonical names.
+    view = _normalize_view(view)
+
+    # Determine whether the horizontal axis should be negated to match
+    # the camera direction that plot_surf uses for the same view string.
+    #   anterior / front / frontal → always negate x (hemisphere-independent)
+    #   medial + rh                → negate x
+    #   lateral + lh               → negate x
+    _negate_x = False
+    if _original_view in ("anterior", "front", "frontal", "bottom", "ventral", "inferior"):
+        _negate_x = True
+    elif hemi is not None:
+        _h = hemi.lower().strip()
+        _is_rh = _h.startswith("r")
+        if _original_view == "medial" and _is_rh:
+            _negate_x = True
+        elif _original_view == "lateral" and not _is_rh:
+            _negate_x = True
+
+    if outline and "color" in kwargs:
+        outline_kwargs["color"] = kwargs.pop("color")
+    if outline and "linewidth" in kwargs:
+        outline_kwargs["linewidth"] = kwargs.pop("linewidth")
+    if outline and "alpha" in kwargs:
+        outline_kwargs["alpha"] = kwargs.pop("alpha")
+
+    # Extract vmin/vmax from kwargs to allow user overrides
+    user_vmin = kwargs.pop("vmin", None)
+    user_vmax = kwargs.pop("vmax", None)
+
+    if isinstance(img, (str, os.PathLike)):
+        img = nib.load(img)
+    if ax is None:
+        fig, ax = plt.subplots()
+        return_fig = True
+    else:
+        return_fig = False
+
+    data = img.get_fdata()
+
+    contains_negative = np.any(data < 0) or diverging
+
+    data = np.where(data == 0, np.nan, data)
+
+    x_ras, y_ras, z_ras = get_world_coords(img)
+
+    if center_coordinate is not None:
+        x_slice_index, y_slice_index, z_slice_index = [
+            int(round(c))
+            for c in nib.affines.apply_affine(
+                np.linalg.inv(img.affine), center_coordinate
+            )
+        ]
+    else:
+        x_slice_index, y_slice_index, z_slice_index = get_center_of_mass(img)
+
+    if view.lower() in ["sagittal", "lateral"]:
+        xlabel = r"P $\rightarrow$ A"
+        ylabel = r"I $\rightarrow$ S"
+        if slice_mm is not None:
+            sl = int(
+                round(
+                    nib.affines.apply_affine(
+                        np.linalg.inv(img.affine), [slice_mm, 0, 0]
+                    )[0]
+                )
+            )
+        if sl is None:
+            sl = x_slice_index
+        data_slice = data[sl, :, :].T
+        ras_extent = (y_ras, z_ras)
+
+    elif view.lower() in ["coronal", "frontal"]:
+        xlabel = r"L $\rightarrow$ R"
+        ylabel = r"I $\rightarrow$ S"
+
+        if slice_mm is not None:
+            sl = int(
+                round(
+                    nib.affines.apply_affine(
+                        np.linalg.inv(img.affine), [0, slice_mm, 0]
+                    )[1]
+                )
+            )
+        if sl is None:
+            sl = y_slice_index
+        data_slice = data[:, sl, :].T
+        ras_extent = (x_ras, z_ras)
+
+    elif view.lower() in ["horizontal", "axial", "transverse"]:
+        xlabel = r"L $\rightarrow$ R"
+        ylabel = r"P $\rightarrow$ A"
+        if slice_mm is not None:
+            sl = int(
+                round(
+                    nib.affines.apply_affine(
+                        np.linalg.inv(img.affine), [0, 0, slice_mm]
+                    )[2]
+                )
+            )
+        if sl is None:
+            sl = z_slice_index
+        data_slice = data[:, :, sl].T
+        ras_extent = (x_ras, y_ras)
+
+    if intensity_projection:
+        mip_axis = (
+            0
+            if view.lower() in ["sagittal", "lateral"]
+            else 1 if view.lower() in ["coronal", "frontal"] else 2
+        )
+        # Use preprocessed data (zeros -> NaNs) so background stays transparent
+        # and extents remain consistent for MIP across views.
+        data_full = data
+        masked_full = ma.masked_invalid(data_full)
+        if intensity_projection == "max":
+            data_slice = masked_full.max(axis=mip_axis).filled(np.nan).T
+        elif intensity_projection == "min":
+            data_slice = masked_full.min(axis=mip_axis).filled(np.nan).T
+        # elif intensity_projection == 'minmax':
+        #     data_slice = masked_full.max(axis=mip_axis).filled(np.nan).T if
+        elif intensity_projection == "mean":
+            data_slice = masked_full.mean(axis=mip_axis).filled(np.nan).T
+        elif intensity_projection == "median":
+            data_slice = ma.median(masked_full, axis=mip_axis).filled(np.nan).T
+        elif intensity_projection == "sum":
+            data_slice = masked_full.sum(axis=mip_axis).filled(np.nan).T
+        elif intensity_projection == "absmax":
+            max_arr = masked_full.max(axis=mip_axis).filled(np.nan)
+            min_arr = masked_full.min(axis=mip_axis).filled(np.nan)
+            pos_abs = np.abs(max_arr)
+            neg_abs = np.abs(min_arr)
+            data_slice = np.where(neg_abs > pos_abs, min_arr, max_arr).T
+        else:
+            raise ValueError(
+                f"Unsupported intensity_projection: {intensity_projection}"
+            )
+
+    # Hide one hemisphere by masking x < 0 (left) or x > 0 (right) in the 2-D slice.
+    # Only meaningful for coronal and axial views where ras_extent[0] is x_ras.
+    if hide_hemi is not None and view.lower() not in ["sagittal", "lateral"]:
+        x_coords = ras_extent[0]
+        if hide_hemi.lower() == "left":
+            mask_cols = x_coords < 0
+        elif hide_hemi.lower() == "right":
+            mask_cols = x_coords > 0
+        else:
+            raise ValueError(f"hide_hemi must be None, 'left', or 'right', got {hide_hemi!r}")
+        data_slice = np.array(data_slice, dtype=float)
+        data_slice[:, mask_cols] = np.nan
+
+    # Store rotation for later application via transform (not data manipulation)
+    rotation_angle = rotate
+
+    data_slice = np.where(abs(data_slice) <= threshold, np.nan, data_slice)
+    data_slice = ma.masked_where(np.isnan(data_slice), data_slice)
+
+    # Determine valid rows/cols based on NaN mask to preserve real-world extents
+    data_slice_filled = data_slice.filled(np.nan)
+    valid_rows = ~np.all(np.isnan(data_slice_filled), axis=1)
+    valid_cols = ~np.all(np.isnan(data_slice_filled), axis=0)
+
+    # Crop the data slice to remove fully-NaN rows and columns
+    data_slice_cropped = data_slice[valid_rows][:, valid_cols]
+    cropped_filled = data_slice_cropped.filled(np.nan)
+
+    # Update extent based on the new bounds; if slice is fully empty, use full world extents
+    col_indices = np.where(valid_cols)[0]
+    row_indices = np.where(valid_rows)[0]
+
+    if col_indices.size == 0 or row_indices.size == 0:
+        extent = [
+            ras_extent[0][0],
+            ras_extent[0][-1],
+            ras_extent[1][0],
+            ras_extent[1][-1],
+        ]
+        data_to_plot = ma.masked_all((2, 2), dtype=float)  # transparent placeholder
+    else:
+        extent = [
+            ras_extent[0][col_indices[0]],
+            ras_extent[0][col_indices[-1]],
+            ras_extent[1][row_indices[0]],
+            ras_extent[1][row_indices[-1]],
+        ]
+        data_to_plot = data_slice_cropped
+
+    # Negate horizontal extent so the slice aligns with plot_surf's
+    # world-coordinate projection for directional views.
+    if _negate_x:
+        extent[0], extent[1] = -extent[0], -extent[1]
+
+    if outline:
+        # Use the cropped slice shape consistently when mapping contour
+        # vertices back to world coordinates, so that outlines align with
+        # the displayed image region.
+        outline_data = np.fliplr(cropped_filled) if flip_lr else cropped_filled
+        contours = find_contours(np.nan_to_num(outline_data), level=0.5)
+        h, w = outline_data.shape
+        interp_x = np.linspace(extent[0], extent[1], w)
+        interp_y = np.linspace(extent[2], extent[3], h)
+
+        # Apply rotation transform if specified
+        if rotation_angle != 0:
+            from matplotlib.transforms import Affine2D
+
+            rotation_transform = Affine2D().rotate_deg(rotation_angle) + ax.transData
+        else:
+            rotation_transform = ax.transData
+
+        for contour in contours:
+            real_world_contour = np.column_stack(
+                [
+                    np.interp(contour[:, 1], np.arange(w), interp_x),
+                    np.interp(contour[:, 0], np.arange(h), interp_y),
+                ]
+            )
+            ax.plot(
+                real_world_contour[:, 0],
+                real_world_contour[:, 1],
+                transform=rotation_transform,
+                **outline_kwargs,
+            )
+        mappable = None
+    else:
+        # Colormap scaling: if the data contains negatives (or diverging=True),
+        # use a symmetric range around zero so post-projection negative values remain visible.
+        # Otherwise, start at the threshold (if provided) for a cleaner lower bound.
+
+        if isinstance(cmap, str):
+            cmap = Colormap(cmap)
+
+        if "norm" in kwargs:
+            vmin = kwargs.pop("vmin", None)
+            vmax = kwargs.pop("vmax", None)
+        else:
+            if contains_negative:
+                vmax = (
+                    np.nanmax(np.abs(cropped_filled))
+                    if col_indices.size and row_indices.size
+                    else 1.0
+                )
+                vmin = -vmax
+            else:
+                if col_indices.size and row_indices.size:
+                    vmin = threshold if threshold != 0 else np.nanmin(cropped_filled)
+                    vmax = np.nanmax(cropped_filled)
+                else:
+                    vmin, vmax = None, None
+
+            # Allow user-provided vmin/vmax to override computed values
+            if user_vmin is not None:
+                vmin = user_vmin
+            if user_vmax is not None:
+                vmax = user_vmax
+            kwargs["vmin"] = vmin
+            kwargs["vmax"] = vmax
+
+        # Apply rotation transform if specified
+        if rotation_angle != 0:
+            from matplotlib.transforms import Affine2D
+
+            rotation_transform = Affine2D().rotate_deg(rotation_angle) + ax.transData
+        else:
+            rotation_transform = ax.transData
+
+        # Flip data left-right for radiological convention if requested
+        if flip_lr:
+            data_to_plot = np.fliplr(data_to_plot)
+
+        mappable = ax.imshow(
+            data_to_plot,
+            cmap=cmap,
+            interpolation="none",
+            alpha=alpha,
+            origin="lower",
+            extent=extent,
+            transform=rotation_transform,
+            **kwargs,
+        )
+
+    # Ensure axes start with negative values if possible
+    if set_axis_extent:
+        # Calculate rotated extent if rotation is applied
+        if rotation_angle != 0:
+            # Get corners of extent rectangle
+            corners = np.array(
+                [
+                    [extent[0], extent[2]],  # bottom-left
+                    [extent[1], extent[2]],  # bottom-right
+                    [extent[1], extent[3]],  # top-right
+                    [extent[0], extent[3]],  # top-left
+                ]
+            )
+            # Rotate corners around origin
+            angle_rad = np.deg2rad(rotation_angle)
+            cos_a, sin_a = np.cos(angle_rad), np.sin(angle_rad)
+            rotated_corners = np.column_stack(
+                [
+                    corners[:, 0] * cos_a - corners[:, 1] * sin_a,
+                    corners[:, 0] * sin_a + corners[:, 1] * cos_a,
+                ]
+            )
+            # Get bounding box of rotated corners
+            xlim = (rotated_corners[:, 0].min(), rotated_corners[:, 0].max())
+            ylim = (rotated_corners[:, 1].min(), rotated_corners[:, 1].max())
+        else:
+            xlim = (min(extent[0], extent[1]), max(extent[0], extent[1]))
+            ylim = (min(extent[2], extent[3]), max(extent[2], extent[3]))
+
+        # Avoid matplotlib warning when min == max by adding small epsilon
+        if xlim[0] == xlim[1]:
+            xlim = (xlim[0] - 0.5, xlim[1] + 0.5)
+        if ylim[0] == ylim[1]:
+            ylim = (ylim[0] - 0.5, ylim[1] + 0.5)
+
+        ax.set_xlim(xlim[0], xlim[1])
+        ax.set_ylim(ylim[0], ylim[1])
+
+    if plot_cut_coord:
+
+        ax.text(
+            0.95,
+            0.95,
+            f"{'xyz'[get_view_index(view)]} = {sl}",
+            transform=ax.transAxes,
+            ha="right",
+        )
+
+    # Optional annotation of the slice position in real-world millimeter space
+    if annotate_slice_mm:
+        # Determine world-space slice coordinate if not explicitly provided
+        world_arrays = [x_ras, y_ras, z_ras]
+        axis_idx = get_view_index(view)
+        axis_char = "xyz"[axis_idx]
+        if slice_mm is not None:
+            slice_mm_value = slice_mm
+        else:
+            # Derive from index -> mm using precomputed world arrays
+            slice_mm_value = world_arrays[axis_idx][sl]
+        ax.text(
+            0.95,
+            0.90,  # slightly lower to avoid overlap with plot_cut_coord text
+            f"{axis_char} = {slice_mm_value:.1f} mm",
+            transform=ax.transAxes,
+            ha="right",
+            va="top",
+        )
+
+    ax.set_aspect("equal")
+
+    # Handle radiological convention: update labels only (no axis inversion)
+    if flip_lr:
+        if view.lower() in ["coronal", "frontal", "horizontal", "axial", "transverse"]:
+            xlabel = r"R $\rightarrow$ L"
+        if view.lower() in ["sagittal", "lateral"]:
+            xlabel = r"A $\rightarrow$ P"  # flip anterior-posterior for sagittal
+
+    # Orientation labels
+    if orientation_labels == "axes":
+        ax.set_xlabel(xlabel)
+        ax.set_ylabel(ylabel)
+    elif orientation_labels == "corner":
+        # Add corner orientation indicators (like clinical neuroradiology images)
+        # Determine labels based on view and flip_lr
+        if view.lower() in ["coronal", "frontal"]:
+            left_label, right_label = ("R", "L") if flip_lr else ("L", "R")
+            top_label, bottom_label = "S", "I"
+        elif view.lower() in ["sagittal", "lateral"]:
+            left_label, right_label = ("A", "P") if flip_lr else ("P", "A")
+            top_label, bottom_label = "S", "I"
+        elif view.lower() in ["horizontal", "axial", "transverse"]:
+            left_label, right_label = ("R", "L") if flip_lr else ("L", "R")
+            top_label, bottom_label = "A", "P"
+        else:
+            left_label, right_label, top_label, bottom_label = "", "", "", ""
+
+        # Add labels outside the axes using data coordinates
+        from matplotlib.patheffects import withStroke
+        xlim = ax.get_xlim()
+        ylim = ax.get_ylim()
+        x_margin = (xlim[1] - xlim[0]) * 0.02
+        y_margin = (ylim[1] - ylim[0]) * 0.02
+
+        text_kwargs = dict(
+            fontsize=12, fontweight="bold", color="white", clip_on=False,
+            path_effects=[withStroke(linewidth=2, foreground="black")]
+        )
+        # Left/right at vertical center, just inside edges
+        ax.text(xlim[0] + x_margin, (ylim[0] + ylim[1]) / 2, left_label,
+                ha="left", va="center", **text_kwargs)
+        ax.text(xlim[1] - x_margin, (ylim[0] + ylim[1]) / 2, right_label,
+                ha="right", va="center", **text_kwargs)
+        # Top/bottom at horizontal center, just inside edges
+        ax.text((xlim[0] + xlim[1]) / 2, ylim[1] - y_margin, top_label,
+                ha="center", va="top", **text_kwargs)
+        ax.text((xlim[0] + xlim[1]) / 2, ylim[0] + y_margin, bottom_label,
+                ha="center", va="bottom", **text_kwargs)
+
+    if colorbar and mappable is not None:
+        cbar = plt.colorbar(mappable, ax=ax, shrink=0.8)
+
+    if return_fig:
+        plt.close()
+        return fig
+
+    return mappable
+
+
+# Backwards compatibility alias
+def plot_img(*args, **kwargs):
+    """Deprecated: use plot_slice instead."""
+    return plot_slice(*args, **kwargs)
+
+
+def translate_com(source_img, target_img, view):
+    com = get_center_of_mass(source_img)
+    com_realworld = nib.affines.apply_affine(source_img.affine, com)
+    com_bg = nib.affines.apply_affine(np.linalg.inv(target_img.affine), com_realworld)
+    sl_bg = int(round(com_bg[get_view_index(view)]))
+    return sl_bg
+
+
+def get_com_view(img, view):
+    if isinstance(img, str) and os.path.isfile(img):
+        img = nib.load(img)
+    com = get_center_of_mass(img)
+    return com[get_view_index(view)]
+
+
+def plot_img_on_bg(
+    img, bg_img, ax=None, view="sagittal", cmap="viridis", zoom=1, **kwargs
+):
+
+    if ax is None:
+        fig, ax = plt.subplots()
+        return_fig = True
+    else:
+        return_fig = False
+
+    com = get_center_of_mass(img)
+    sl = int(round(com[get_view_index(view)]))
+    com_realworld = nib.affines.apply_affine(img.affine, com)
+    com_bg = nib.affines.apply_affine(np.linalg.inv(bg_img.affine), com_realworld)
+    sl_bg = int(round(com_bg[get_view_index(view)]))
+
+    plot_slice(bg_img, ax=ax, cmap="gray", view=view, sl=sl_bg)
+    xlim = ax.get_xlim()
+    ylim = ax.get_ylim()
+
+    if not isinstance(img, list):
+        img = [img]
+
+    for iimg in img:
+        mappable = plot_slice(iimg, ax=ax, view=view, cmap=cmap, sl=sl, **kwargs)
+
+    ax.set_xlim([l / zoom for l in xlim])
+    ax.set_ylim([l / zoom for l in ylim])
+
+    if return_fig:
+        plt.close()
+        return fig
+    else:
+        return mappable
+
+
+def _get_object_bounds(json_content, transform):
+    xmin, xmax, ymin, ymax = json_content["metadata"]["bounds"]
+    x0, y0 = transform.transform((xmin, ymin))
+    x1, y1 = transform.transform((xmax, ymax))
+
+    xmin, xmax = min(x0, x1), max(x0, x1)
+    ymin, ymax = min(y0, y1), max(y0, y1)
+
+    # A combination of a proportional factor (fraction of the drawing)
+    # and a guestimate of the linewidth
+    xmargin = (xmax - xmin) * 0.025
+    ymargin = (ymax - ymin) * 0.025
+    return xmin, xmax, ymin, ymax
+
+
+def plot_brain_schematics(ax, view="sagittal", **kwargs):
+    from nilearn.plotting import glass_brain
+
+    if view.lower() == "sagittal":
+        direction = "x"
+        shift_x = -5
+        shift_y = -5
+    elif view.lower() == "coronal":
+        direction = "y"
+        shift_x = 5
+        shift_y = -5
+    elif view.lower() == "horizontal":
+        direction = "z"
+        shift_x = 5
+        shift_y = 5
+    shift_x = 0
+    shift_y = 0
+
+    get_axis_bg_color = ax.get_facecolor()
+
+    black_bg = colors.colorConverter.to_rgba(
+        get_axis_bg_color
+    ) == colors.colorConverter.to_rgba("k")
+
+    json_filename, transform = glass_brain._get_json_and_transform(direction)
+    with open(json_filename) as json_file:
+        json_content = json.load(json_file)
+
+    zoom = 1
+
+    xmin, xmax, ymin, ymax = json_content["metadata"]["bounds"]
+    cx = (xmin + xmax) / 2
+    cy = (ymin + ymax) / 2
+
+    # Compose zoom-around-center transform
+    transform = (
+        transforms.Affine2D()
+        .translate(-cx, -cy)  # move to center
+        .scale(zoom, zoom)  # zoom
+        .translate(cx, cy)  # move back
+        .translate(shift_x, shift_y)  # final shift
+        + transform
+    )
+    # transform = transform.scale(zoom, zoom)
+    linewidth_scale = kwargs.pop("linewidth", None)
+    abs_linewidth = kwargs.pop("abs_linewidth", None)
+    mpl_patches = glass_brain._get_mpl_patches(
+        json_content,
+        transform=transform + ax.transData,
+        invert_color=black_bg,
+        **kwargs,
+    )
+
+    # Build absolute linewidth mapping: original widths -> replacement widths
+    if abs_linewidth is not None:
+        original_widths = sorted({p["linewidth"] for p in json_content["paths"]})
+        if isinstance(abs_linewidth, (int, float)):
+            lw_map = {ow: abs_linewidth for ow in original_widths}
+        else:
+            # list/tuple: map to [fine, medium, bold] in ascending order
+            abs_linewidth = list(abs_linewidth)
+            if len(abs_linewidth) != len(original_widths):
+                raise ValueError(
+                    f"abs_linewidth list must have {len(original_widths)} values "
+                    f"(one per width category: {original_widths}), "
+                    f"got {len(abs_linewidth)}"
+                )
+            lw_map = dict(zip(original_widths, abs_linewidth))
+
+    for mpl_patch in mpl_patches:
+        if abs_linewidth is not None:
+            # Map original width to user-specified absolute width
+            orig_lw = mpl_patch.get_linewidth()
+            mpl_patch.set_linewidth(lw_map.get(orig_lw, orig_lw))
+        elif linewidth_scale is not None:
+            mpl_patch.set_linewidth(mpl_patch.get_linewidth() * linewidth_scale)
+        ax.add_patch(mpl_patch)
+
+    bounds = _get_object_bounds(json_content, transform)
+    ax.set_xlim(bounds[0], bounds[1])
+    ax.set_ylim(bounds[2], bounds[3])
+    ax.set_aspect("equal")
+    return bounds
+
+
+def get_world_extent(img):
+    from itertools import product
+
+    shape = np.array(img.shape[:3])
+    corners = np.array(
+        list(product([0, shape[0] - 1], [0, shape[1] - 1], [0, shape[2] - 1]))
+    )
+    real = nib.affines.apply_affine(img.affine, corners)
+    return np.min(real, axis=0), np.max(real, axis=0)
+
+
+def plot_image_extent(
+    img,
+    ax,
+    view="sagittal",
+    color="r",
+    zoom_in=False,
+    pad=5,
+    slice_mm=None,
+    threshold=0,
+    hemi=None,
+    **kwargs,
+):
+    if isinstance(img, (str, os.PathLike)):
+        img = nib.load(img)
+
+    # Determine whether the horizontal axis must be negated to match
+    # the coordinate convention used by plot_slice / plot_surf.
+    _original_view = view.lower().strip() if isinstance(view, str) else str(view)
+    _negate_x = _original_view in ("anterior", "front", "frontal",
+                                    "bottom", "ventral", "inferior")
+    if not _negate_x and hemi is not None:
+        _h = hemi.lower().strip()
+        _is_rh = _h.startswith("r")
+        if _original_view == "medial" and _is_rh:
+            _negate_x = True
+        elif _original_view == "lateral" and not _is_rh:
+            _negate_x = True
+
+    # Use get_plane_indices which normalises view aliases internally.
+    axes = get_plane_indices(view)
+
+    if slice_mm is not None:
+        # Get extent at specific slice
+        data = img.get_fdata()
+        slice_axis = get_view_index(view)
+
+        # Convert slice_mm to voxel index
+        slice_voxel = int(
+            round(
+                nib.affines.apply_affine(
+                    np.linalg.inv(img.affine),
+                    [
+                        slice_mm if slice_axis == 0 else 0,
+                        slice_mm if slice_axis == 1 else 0,
+                        slice_mm if slice_axis == 2 else 0,
+                    ],
+                )[slice_axis]
+            )
+        )
+
+        # Extract the slice
+        if slice_axis == 0:
+            slice_data = data[slice_voxel, :, :]
+        elif slice_axis == 1:
+            slice_data = data[:, slice_voxel, :]
+        else:
+            slice_data = data[:, :, slice_voxel]
+
+        # Find extent above threshold in the slice
+        nonzero_coords = np.argwhere(np.abs(slice_data) > threshold)
+        if len(nonzero_coords) == 0:
+            return  # Nothing to plot
+
+        # Get min/max in slice coordinates
+        min_coords = nonzero_coords.min(axis=0)
+        max_coords = nonzero_coords.max(axis=0)
+
+        # Convert to voxel coordinates in 3D
+        if slice_axis == 0:  # sagittal
+            voxel_min = np.array([slice_voxel, min_coords[0], min_coords[1]])
+            voxel_max = np.array([slice_voxel, max_coords[0], max_coords[1]])
+        elif slice_axis == 1:  # coronal
+            voxel_min = np.array([min_coords[0], slice_voxel, min_coords[1]])
+            voxel_max = np.array([max_coords[0], slice_voxel, max_coords[1]])
+        else:  # horizontal
+            voxel_min = np.array([min_coords[0], min_coords[1], slice_voxel])
+            voxel_max = np.array([max_coords[0], max_coords[1], slice_voxel])
+
+        # Convert to world coordinates
+        world_min = nib.affines.apply_affine(img.affine, voxel_min)
+        world_max = nib.affines.apply_affine(img.affine, voxel_max)
+
+        x_min = world_min[axes[0]]
+        x_max = world_max[axes[0]]
+        y0 = world_min[axes[1]]
+        height = world_max[axes[1]] - world_min[axes[1]]
+    else:
+        # Compute extent from voxels above *threshold* (projected along
+        # the slice axis).  Falls back to full volume extent when no
+        # voxels survive the threshold.
+        data = img.get_fdata()
+        above = np.argwhere(np.abs(data) > threshold)
+        if above.size == 0:
+            # Nothing above threshold – use full volume corners
+            world_extent = get_world_extent(img)
+            x_min = world_extent[0][axes[0]]
+            x_max = world_extent[1][axes[0]]
+            y0 = world_extent[0][axes[1]]
+            height = world_extent[1][axes[1]] - world_extent[0][axes[1]]
+        else:
+            voxel_min = above.min(axis=0)
+            voxel_max = above.max(axis=0)
+            world_min = nib.affines.apply_affine(img.affine, voxel_min)
+            world_max = nib.affines.apply_affine(img.affine, voxel_max)
+            x_min = world_min[axes[0]]
+            x_max = world_max[axes[0]]
+            y0 = world_min[axes[1]]
+            height = world_max[axes[1]] - world_min[axes[1]]
+
+    # Negate horizontal extent for views where plot_slice / plot_surf
+    # flip the x-axis (e.g. frontal, bottom).
+    if _negate_x:
+        x_min, x_max = -x_max, -x_min
+
+    x0 = x_min
+    width = x_max - x_min
+
+    if "facecolor" not in kwargs:
+        kwargs["facecolor"] = "none"
+    if "edgecolor" not in kwargs:
+        kwargs["edgecolor"] = color
+    rect = patches.Rectangle((x0, y0), width, height, **kwargs)
+    ax.add_patch(rect)
+
+    if zoom_in:
+        ax.set_xlim(x0 - pad, x0 + width + pad)
+        ax.set_ylim(y0 - pad, y0 + height + pad)
+
+
+def apply_brain_mask(
+    img,
+    template="MNI152NLin2009cAsym",
+    resolution=1,
+    desc="brain",
+    suffix="mask",
+):
+    """
+    Mask a NIfTI image in MNI space with a brain mask from TemplateFlow.
+
+    The brain mask is resampled to match the input image grid using
+    nearest-neighbour interpolation and then applied multiplicatively.
+
+    Parameters
+    ----------
+    img : str, PathLike, or Nifti1Image
+        NIfTI image (or path) in MNI space to be masked.
+    template : str, optional
+        TemplateFlow template name (default ``"MNI152NLin2009cAsym"``).
+    resolution : int, optional
+        Resolution index passed to ``templateflow.api.get``
+        (default ``1``).
+    desc : str, optional
+        Description label for the mask (default ``"brain"``).
+    suffix : str, optional
+        BIDS suffix for the mask (default ``"mask"``).
+
+    Returns
+    -------
+    masked_img : Nifti1Image
+        Brain-masked image with the same affine, header, and data type
+        as the input.
+    """
+    from templateflow import api as tflow
+    from nilearn.image import resample_to_img
+
+    if isinstance(img, (str, os.PathLike)):
+        img = nib.load(img)
+
+    mask_path = tflow.get(
+        template,
+        resolution=resolution,
+        desc=desc,
+        suffix=suffix,
+    )
+    mask_img = nib.load(mask_path)
+
+    mask_resampled = resample_to_img(
+        mask_img, img, interpolation="nearest",
+        force_resample=True, copy_header=True,
+    )
+
+    data = img.get_fdata()
+    mask_data = mask_resampled.get_fdata().astype(bool)
+    masked_data = data * mask_data
+
+    return nib.Nifti1Image(masked_data, img.affine, img.header)