visualize.py

"""Generate qualitative figures: input / ground truth / prediction panels.

Usage::

    python scripts/visualize.py --checkpoint artifacts/phase0_unet2d/best.pt --n 8
    python scripts/visualize.py --checkpoint artifacts/phase0_unet2d/best.pt --worst

Quantitative tables hide failure modes. ``--worst`` selects the lowest-Dice
cases in the split rather than a random sample, because those are the ones worth
looking at: a figure of eight easy cases proves nothing that the mean Dice did
not already say.
"""

from __future__ import annotations

import argparse
import logging
import sys
from pathlib import Path

import numpy as np
import torch

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

import matplotlib  # noqa: E402

matplotlib.use("Agg")  # no display in CI or over SSH
import matplotlib.pyplot as plt  # noqa: E402

from src.data.datasets import build_cheng_datasets  # noqa: E402
from src.detection.bbox import masks_to_boxes  # noqa: E402
from src.eval.metrics import dice_score  # noqa: E402
from src.models.registry import build_model  # noqa: E402

logger = logging.getLogger("visualize")


def load_model(checkpoint_path: Path, device: torch.device):
    checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=False)
    config = checkpoint["config"]
    model_config = dict(config["model"])
    name = model_config.pop("name")
    model = build_model(name, **model_config)
    model.load_state_dict(checkpoint["model_state"])
    model.to(device).eval()
    return model, config, checkpoint.get("metrics", {})


@torch.no_grad()
def collect_predictions(model, dataset, device, threshold: float):
    """Predict every case in the dataset, returning image/target/prediction/Dice."""
    records = []
    loader = torch.utils.data.DataLoader(dataset, batch_size=16, shuffle=False)

    for batch in loader:
        images = batch["image"].to(device)
        probabilities = torch.sigmoid(model(images).float()).cpu().numpy()
        predictions = probabilities > threshold

        for i in range(images.size(0)):
            image = batch["image"][i, 0].numpy()
            target = batch["mask"][i, 0].numpy() > 0.5
            prediction = predictions[i, 0]
            records.append(
                {
                    "image": image,
                    "target": target,
                    "prediction": prediction,
                    "dice": dice_score(prediction, target),
                    "index": int(batch["index"][i]),
                }
            )
    return records


def plot_panel(records: list[dict], out_path: Path, title: str) -> None:
    """Four columns per case: image, ground truth, prediction, overlay + boxes."""
    n = len(records)
    fig, axes = plt.subplots(n, 4, figsize=(11, 2.8 * n), squeeze=False)

    for row, record in enumerate(records):
        image = record["image"]
        target, prediction = record["target"], record["prediction"]

        axes[row][0].imshow(image, cmap="gray")
        axes[row][0].set_ylabel(f"#{record['index']}\nDice {record['dice']:.3f}", fontsize=9)

        axes[row][1].imshow(target, cmap="gray")
        axes[row][2].imshow(prediction, cmap="gray")

        # Overlay: green = correct, red = missed (FN), blue = spurious (FP).
        # Colouring the two error types differently distinguishes a model that
        # under-segments from one that over-segments, which a single-colour
        # overlay cannot.
        overlay = np.stack([_norm(image)] * 3, axis=-1)
        overlay[target & prediction] = [0.1, 0.9, 0.2]
        overlay[target & ~prediction] = [0.95, 0.15, 0.15]
        overlay[~target & prediction] = [0.2, 0.45, 1.0]
        axes[row][3].imshow(overlay)

        for box in masks_to_boxes(prediction.astype(np.uint8), min_area=16):
            axes[row][3].add_patch(
                plt.Rectangle(
                    (box.x1, box.y1),
                    box.width,
                    box.height,
                    fill=False,
                    edgecolor="yellow",
                    linewidth=1.2,
                )
            )

        for col, label in enumerate(["input", "ground truth", "prediction", "overlay + boxes"]):
            axes[row][col].set_xticks([])
            axes[row][col].set_yticks([])
            if row == 0:
                axes[row][col].set_title(label, fontsize=10)

    fig.suptitle(title, fontsize=12)
    fig.text(
        0.5,
        0.005,
        "overlay: green = true positive, red = false negative, blue = false positive",
        ha="center",
        fontsize=8,
        color="#555",
    )
    fig.tight_layout(rect=[0, 0.015, 1, 0.985])
    out_path.parent.mkdir(parents=True, exist_ok=True)
    fig.savefig(out_path, dpi=130, bbox_inches="tight")
    plt.close(fig)


def _norm(image: np.ndarray) -> np.ndarray:
    lo, hi = float(image.min()), float(image.max())
    if hi - lo < 1e-8:
        return np.zeros_like(image)
    return (image - lo) / (hi - lo)


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description="Save qualitative prediction figures")
    parser.add_argument("--checkpoint", type=Path, default=Path("artifacts/phase0_unet2d/best.pt"))
    parser.add_argument("--cache", type=Path, default=Path("data/processed/cheng_256.npz"))
    parser.add_argument("--out-dir", type=Path, default=Path("reports/figures"))
    parser.add_argument("--split", default="test", choices=["train", "val", "test"])
    parser.add_argument("--n", type=int, default=6)
    parser.add_argument("--worst", action="store_true", help="show the lowest-Dice cases")
    parser.add_argument("--device", default="cuda")
    parser.add_argument("--seed", type=int, default=42)
    args = parser.parse_args(argv)

    logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s")

    if not args.checkpoint.exists():
        logger.error("checkpoint not found: %s", args.checkpoint)
        return 1

    device = torch.device(args.device if torch.cuda.is_available() else "cpu")
    model, config, metrics = load_model(args.checkpoint, device)

    datasets, _ = build_cheng_datasets(args.cache, seed=args.seed, augment_train=False)
    dataset = datasets[args.split]

    logger.info("predicting %d cases from the %s split", len(dataset), args.split)
    records = collect_predictions(model, dataset, device, config.get("threshold", 0.5))
    records.sort(key=lambda r: r["dice"])

    dices = np.array([r["dice"] for r in records])
    logger.info(
        "Dice over the split: mean %.4f, median %.4f, min %.4f, max %.4f",
        dices.mean(),
        np.median(dices),
        dices.min(),
        dices.max(),
    )

    run_name = args.checkpoint.parent.name

    worst = records[: args.n]
    plot_panel(
        worst,
        args.out_dir / f"{run_name}_{args.split}_worst.png",
        f"{run_name} — {args.n} worst cases ({args.split} split)",
    )

    if not args.worst:
        # An evenly spaced sample across the sorted Dice range, so the figure
        # spans the full spread of performance rather than only its extremes.
        idx = np.linspace(0, len(records) - 1, args.n).round().astype(int)
        spread = [records[i] for i in np.unique(idx)]
        plot_panel(
            spread,
            args.out_dir / f"{run_name}_{args.split}_spread.png",
            f"{run_name} — cases spanning the Dice range ({args.split} split)",
        )

        best = records[-args.n :][::-1]
        plot_panel(
            best,
            args.out_dir / f"{run_name}_{args.split}_best.png",
            f"{run_name} — {args.n} best cases ({args.split} split)",
        )

    # Dice histogram: the shape of the distribution says more than its mean.
    fig, ax = plt.subplots(figsize=(6, 3.2))
    ax.hist(dices, bins=30, color="#4C78A8", edgecolor="white")
    ax.axvline(dices.mean(), color="#E45756", linestyle="--", label=f"mean {dices.mean():.3f}")
    ax.set_xlabel("Dice")
    ax.set_ylabel("cases")
    ax.set_title(f"{run_name} — Dice distribution ({args.split})")
    ax.legend()
    fig.tight_layout()
    fig.savefig(args.out_dir / f"{run_name}_{args.split}_dice_hist.png", dpi=130)
    plt.close(fig)

    logger.info("figures written to %s", args.out_dir)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())