bundle_qa.py

#!/usr/bin/env python3
"""
QA rendering of bundle segmentations with glass-brain overlay.

Renders sagittal, coronal, and axial maximum-intensity projections of
each bundle (values 1-N colored via jet colormap) with a transparent
glass-brain outline computed from the brain mask.

Glass brain: upsample mask -> smooth -> threshold -> dilate -> shell,
then sum-project for an X-ray/outline effect.

Bundle: values 1-N rendered with discrete jet colormap, 0 = transparent.
"""

import argparse
import os
from pathlib import Path

import nibabel as nib
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.colors as mcolors
from matplotlib.colors import ListedColormap
from scipy.ndimage import zoom, gaussian_filter, binary_dilation


# ── Glass-brain shell ──────────────────────────────────────────────────

def _compute_shell(binary_data, vox_sizes, target_vox=0.5, stdev_mm=2.0,
                   dilate_passes=2):
    """Upsample, smooth, threshold and dilate a binary mask, returning
    the outer shell and the zoom factors used."""
    zoom_factors = vox_sizes / target_vox
    data_up = zoom(binary_data.astype(float), zoom_factors, order=1)
    stdev_vox = stdev_mm / target_vox
    data_smooth = gaussian_filter(data_up, sigma=stdev_vox)
    data_thres = (data_smooth >= 0.5).astype(np.uint8)
    data_dilated = binary_dilation(data_thres, iterations=dilate_passes).astype(np.uint8)
    shell = (data_dilated - data_thres).astype(np.uint8)
    return shell, zoom_factors


def compute_glass_brain(mask_path, target_vox=0.5, stdev_mm=2.0,
                        dilate_passes=2):
    """Load a brain-mask NIfTI and return the glass-brain shell volume
    (at upsampled resolution) together with the zoom factors."""
    img = nib.load(mask_path)
    affine = img.affine
    vox_sizes = np.sqrt((affine[:3, :3] ** 2).sum(axis=0))
    binary_data = img.get_fdata() > 0
    shell, zoom_factors = _compute_shell(
        binary_data, vox_sizes, target_vox, stdev_mm, dilate_passes,
    )
    return shell, zoom_factors


# ── Projection helpers ─────────────────────────────────────────────────

def _orient_projection(proj, axis):
    """Transpose and flip a 2D projection for standard radiological display.

    After projecting along *axis* from a (X, Y, Z) volume the remaining
    2-D array has axes (kept-dim-0, kept-dim-1).  We rotate/flip so that
    the superior direction points up and the image looks anatomically
    conventional.
    """
    # axis 0 (sagittal): remaining (Y, Z) -> rows=Z cols=Y, flip rows S-up
    # axis 1 (coronal):  remaining (X, Z) -> rows=Z cols=X, flip rows S-up
    # axis 2 (axial):    remaining (X, Y) -> rows=Y cols=X, flip rows A-up
    proj = proj.T[::-1, :]
    return proj


def project_glass_brain(shell, axis):
    """Sum-projection of the glass-brain shell (gives X-ray outline)."""
    proj = shell.astype(float).sum(axis=axis)
    return _orient_projection(proj, axis)


def project_bundle(bundle_data, axis):
    """Max-intensity projection of bundle labels (0-N) along *axis*."""
    proj = bundle_data.max(axis=axis)
    return _orient_projection(proj, axis)


# ── Discrete colormap for bundle segments (1-N) ───────────────────────

def make_bundle_cmap(n_labels=10):
    """0 -> transparent, 1..n_labels -> evenly-spaced jet colours."""
    jet = plt.cm.jet
    colors = [(0, 0, 0, 0)]  # label 0 = transparent
    for i in range(n_labels):
        r, g, b, _ = jet(i / max(n_labels - 1, 1))
        colors.append((r, g, b, 1.0))
    cmap = ListedColormap(colors, name="bundle_discrete")
    bounds = np.arange(-0.5, n_labels + 1.5, 1)
    norm = mcolors.BoundaryNorm(bounds, cmap.N)
    return cmap, norm


# ── Main rendering function ────────────────────────────────────────────

def render_bundle_qa(mask_path, bundle_path, out_png, title=None,
                     target_vox=0.5, glass_alpha=0.4, n_labels=10,
                     precomputed_shell=None, precomputed_zoom=None):
    """Render sagittal / coronal / axial MIP views of a single bundle
    with a glass-brain outline underneath.

    If *precomputed_shell* and *precomputed_zoom* are supplied, the
    (expensive) glass-brain computation is skipped.
    """

    # ── Glass brain ──
    if precomputed_shell is not None and precomputed_zoom is not None:
        shell = precomputed_shell
        zoom_factors = precomputed_zoom
    else:
        shell, zoom_factors = compute_glass_brain(mask_path,
                                                  target_vox=target_vox)

    # ── Bundle (upsample to match shell resolution) ──
    bimg = nib.load(bundle_path)
    bdata = np.rint(bimg.get_fdata()).astype(np.uint8)
    bdata_up = zoom(bdata.astype(float), zoom_factors, order=0)
    bdata_up = np.rint(bdata_up).astype(np.uint8)

    if title is None:
        name = Path(bundle_path).name
        title = name.replace(".nii.gz", "").replace(".nii", "")

    cmap, norm = make_bundle_cmap(n_labels)

    view_labels = ["Sagittal", "Coronal", "Axial"]
    proj_axes = [0, 1, 2]

    fig, axs = plt.subplots(1, 3, figsize=(15, 5), facecolor="white")
    fig.suptitle(title, color="black", fontsize=16, fontweight="bold")

    for ax, proj_axis, vlabel in zip(axs, proj_axes, view_labels):
        # Glass-brain outline (sum projection, normalised)
        gb_proj = project_glass_brain(shell, proj_axis)
        gb_max = gb_proj.max()
        if gb_max > 0:
            gb_proj = gb_proj / gb_max

        # Bundle MIP
        b_proj = project_bundle(bdata_up, proj_axis)

        # Draw glass brain outline in gray on white background
        ax.imshow(gb_proj, cmap="gray_r", vmin=0, vmax=1,
                  alpha=glass_alpha, aspect="equal")

        # Draw bundle on top (mask zeros for transparency)
        b_masked = np.ma.masked_where(b_proj == 0, b_proj)
        ax.imshow(b_masked, cmap=cmap, norm=norm,
                  aspect="equal", interpolation="nearest")

        ax.set_title(vlabel, color="black", fontsize=11)
        ax.axis("off")

    fig.tight_layout(rect=[0, 0, 1, 0.93])
    os.makedirs(os.path.dirname(out_png) or ".", exist_ok=True)
    fig.savefig(out_png, dpi=200, facecolor="white", bbox_inches="tight")
    plt.close(fig)

    print(f"Saved: {out_png}")


# ── Batch mode: process all bundles in a directory ─────────────────────

def render_all_bundles(mask_path, bundle_dir, out_dir, target_vox=0.5,
                       glass_alpha=0.7, n_labels=10):
    """Render QA PNGs for every *.nii.gz in *bundle_dir*.
    The glass brain is computed once and reused for all bundles."""
    bundle_dir = Path(bundle_dir)
    out_dir = Path(out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)

    bundles = sorted(bundle_dir.glob("*.nii.gz"))
    if not bundles:
        print(f"WARNING: no .nii.gz files found in {bundle_dir}")
        return

    print(f"Computing glass brain from {mask_path} ...")
    shell, zoom_factors = compute_glass_brain(mask_path,
                                              target_vox=target_vox)
    print(f"Glass brain shell shape: {shell.shape}")

    for seg in bundles:
        stem = seg.name.replace(".nii.gz", "").replace(".nii", "")
        out_png = out_dir / f"{stem}.png"
        print(f"Rendering {stem} ...")
        render_bundle_qa(
            mask_path=mask_path,
            bundle_path=str(seg),
            out_png=str(out_png),
            title=stem,
            target_vox=target_vox,
            glass_alpha=glass_alpha,
            n_labels=n_labels,
            precomputed_shell=shell,
            precomputed_zoom=zoom_factors,
        )

    print(f"All QA PNGs saved to: {out_dir}")


# ── CLI ────────────────────────────────────────────────────────────────

def _build_argparser():
    p = argparse.ArgumentParser(
        description="Render sagittal/coronal/axial QA views of bundles "
                    "with glass-brain overlay.",
    )
    # Single-bundle mode
    p.add_argument("--mask", required=True,
                   help="Brain mask NIfTI (binary).")
    p.add_argument("--bundle", default=None,
                   help="Single bundle NIfTI (values 0-N). "
                        "Mutually exclusive with --bundle_dir.")
    p.add_argument("--out", default=None,
                   help="Output PNG path (single-bundle mode).")
    # Batch mode
    p.add_argument("--bundle_dir", default=None,
                   help="Directory of bundle NIfTIs (batch mode).")
    p.add_argument("--out_dir", default=None,
                   help="Output directory for PNGs (batch mode).")
    # Common
    p.add_argument("--title", default=None,
                   help="Title (single-bundle mode; defaults to filename).")
    p.add_argument("--n_labels", type=int, default=10,
                   help="Number of label values in bundles (default 10).")
    p.add_argument("--target_vox", type=float, default=0.5,
                   help="Glass-brain resolution in mm (default 0.5).")
    p.add_argument("--glass_alpha", type=float, default=0.4,
                   help="Glass-brain opacity (default 0.4).")
    return p


def main():
    args = _build_argparser().parse_args()

    if args.bundle_dir:
        # Batch mode
        if not args.out_dir:
            print("ERROR: --out_dir is required in batch mode")
            raise SystemExit(2)
        render_all_bundles(
            mask_path=args.mask,
            bundle_dir=args.bundle_dir,
            out_dir=args.out_dir,
            target_vox=args.target_vox,
            glass_alpha=args.glass_alpha,
            n_labels=args.n_labels,
        )
    elif args.bundle:
        # Single-bundle mode
        out = args.out or "qa.png"
        render_bundle_qa(
            mask_path=args.mask,
            bundle_path=args.bundle,
            out_png=out,
            title=args.title,
            n_labels=args.n_labels,
            target_vox=args.target_vox,
            glass_alpha=args.glass_alpha,
        )
    else:
        print("ERROR: provide either --bundle (single) or "
              "--bundle_dir (batch)")
        raise SystemExit(2)


if __name__ == "__main__":
    main()