make_figures.py

#!/usr/bin/env python3
"""Render example slice figures for each pipeline stage into images/.

One shared axial slice (chosen as the slice carrying the most distinct tissue
labels) is used across all figures so they are directly comparable.

Outputs:
  images/00_inputs.png   - stage 0: T1 + Allen annotation (input volumes)
  images/02_classify.png - stage 2: 9-class segmentation over the T1
  images/03_priors.png   - stage 3: the 8 probabilistic priors
  images/04_coarse.png   - stage 4: the 4-class / 3-class / cbmerge schemes
"""
import os

import numpy as np
import SimpleITK as sitk
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.colors import ListedColormap
from matplotlib.patches import Patch

os.makedirs("images", exist_ok=True)

T1 = "mni_icbm152_t1_tal_nlin_sym_09b_hires.nii.gz"
ANN = "annotation_full.nii.gz"
CLS = "classify.nii.gz"

# 9-class tissue colours (0..8)
TISSUE = [
    ("background", "#000000"), ("CSF", "#3b6ea5"), ("GM", "#9e9e9e"),
    ("WM", "#f0f0f0"), ("ventricle", "#36c9c9"), ("cerebGM", "#e8b23a"),
    ("cerebWM", "#e2742a"), ("brainStem", "#4caf50"), ("deepGM", "#d64550"),
]
TCMAP = ListedColormap([c for _, c in TISSUE])


def arr(path):
    return sitk.GetArrayFromImage(sitk.ReadImage(path))  # [z, y, x]


def sl(a, z):
    return np.rot90(a[z], 2)  # upright axial, rotated a further 90° CCW


def norm(img):
    lo, hi = np.percentile(img, [1, 99])
    return np.clip((img - lo) / (hi - lo + 1e-9), 0, 1)


def pick_slice(lab):
    # axial slice with the most distinct nonzero labels (ties -> most fg voxels)
    best, bz = (-1, -1), 0
    for z in range(lab.shape[0]):
        s = lab[z]
        score = (len(np.unique(s[s > 0])), int((s > 0).sum()))
        if score > best:
            best, bz = score, z
    return bz


def figure_inputs(t1, ann, z):
    fig, ax = plt.subplots(1, 2, figsize=(9, 5.2))
    ax[0].imshow(sl(t1, z), cmap="gray"); ax[0].set_title("MNI ICBM152 T1 (input)")
    # annotation: random colours per id for visual separation
    a = sl(ann, z)
    rng = np.random.default_rng(0)
    ids = np.unique(a)
    lut = {i: rng.random(3) for i in ids}; lut[0] = (0, 0, 0)
    rgb = np.zeros((*a.shape, 3))
    for i in ids:
        rgb[a == i] = lut[i]
    ax[1].imshow(rgb); ax[1].set_title("Allen annotation: structure ids (input)")
    for a_ in ax:
        a_.axis("off")
    fig.suptitle("Stage 0 — inputs", fontsize=13)
    fig.tight_layout(); fig.savefig("images/00_inputs.png", dpi=120, bbox_inches="tight")
    plt.close(fig)


def figure_classify(t1, lab, z):
    fig, ax = plt.subplots(figsize=(6.4, 6.4))
    ax.imshow(sl(t1, z), cmap="gray")
    m = sl(lab, z).astype(float)
    ax.imshow(np.ma.masked_where(m == 0, m), cmap=TCMAP, vmin=0, vmax=8, alpha=0.6)
    ax.axis("off"); ax.set_title("Stage 2 — 9-class segmentation over T1", fontsize=12)
    leg = [Patch(facecolor=c, edgecolor="k", label=f"{i} {n}")
           for i, (n, c) in enumerate(TISSUE) if i > 0]
    ax.legend(handles=leg, loc="center left", bbox_to_anchor=(1.0, 0.5),
              fontsize=8, frameon=False)
    fig.tight_layout(); fig.savefig("images/02_classify.png", dpi=120, bbox_inches="tight")
    plt.close(fig)


def figure_priors(z):
    names = [n for n, _ in TISSUE[1:]]  # 1..8
    fig, ax = plt.subplots(2, 4, figsize=(11, 6))
    for c in range(1, 9):
        a = ax.flat[c - 1]
        p = sl(arr(f"prior{c:02d}.nii.gz"), z)
        a.imshow(p, cmap="magma", vmin=0, vmax=1)
        a.set_title(f"prior{c:02d}  {names[c - 1]}", fontsize=9); a.axis("off")
    fig.suptitle("Stage 3 — probabilistic priors (0-1)", fontsize=13)
    fig.tight_layout(); fig.savefig("images/03_priors.png", dpi=120, bbox_inches="tight")
    plt.close(fig)


def figure_coarse(z):
    schemes = [("classify_4class.nii.gz", "4-class", 4),
               ("classify_3class.nii.gz", "3-class", 3),
               ("classify_cbmerge.nii.gz", "cbmerge", 7)]
    cmap = ListedColormap(["#000000"] + list(plt.get_cmap("tab10").colors))
    fig, ax = plt.subplots(1, 3, figsize=(12, 4.6))
    for a, (f, name, k) in zip(ax, schemes):
        m = sl(arr(f), z).astype(float)
        a.imshow(np.ma.masked_where(m == 0, m), cmap=cmap, vmin=0, vmax=10)
        a.set_title(f"{name} ({k} classes)", fontsize=11); a.axis("off")
    fig.suptitle("Stage 4 — coarser schemes", fontsize=13)
    fig.tight_layout(); fig.savefig("images/04_coarse.png", dpi=120, bbox_inches="tight")
    plt.close(fig)


def main():
    t1, ann, lab = arr(T1), arr(ANN), arr(CLS)
    z = pick_slice(lab)
    print(f"using axial slice z={z} (of {lab.shape[0]})")
    figure_inputs(t1, ann, z)
    figure_classify(t1, lab, z)
    figure_priors(z)
    figure_coarse(z)
    print("wrote images/00_inputs.png 02_classify.png 03_priors.png 04_coarse.png")


if __name__ == "__main__":
    main()