evaluate.py

"""Single benchmarking entrypoint for every model family.

Runs each trained checkpoint over a held-out split, computes the full metric
bundle, and writes a comparison report in markdown, HTML and JSON.

Usage::

    # Benchmark every 2-D checkpoint under artifacts/ on the Cheng test split
    python evaluate.py --task segmentation-2d

    # Add detection, 3-D and interpolation results into one report
    python evaluate.py --task all

    # CI mode: small sample, and fail the build if metrics regressed
    python evaluate.py --task segmentation-2d --subset 64 \
        --gate reports/baseline.json --tolerance 0.02

    # Record the current results as the new baseline
    python evaluate.py --task segmentation-2d --save-baseline reports/baseline.json

``--task all`` reads the JSON summaries that the detection, volumetric and
interpolation training scripts write beside their checkpoints, rather than
re-running them. Each family gets its **own table**: detection mAP,
reconstruction PSNR and segmentation Dice are not comparable quantities, and
one wide table of mostly-empty cells invites exactly the cross-task comparison
that would be meaningless.

The regression gate compares against a stored baseline and exits non-zero if any
metric falls more than ``tolerance`` below it. Higher-is-better and
lower-is-better metrics are handled separately — a *drop* in Dice is a
regression, whereas a drop in Hausdorff distance is an improvement, and treating
them the same way is an easy way to build a gate that passes when it should
fail.
"""

from __future__ import annotations

import argparse
import json
import logging
import sys
import time
from pathlib import Path
from typing import Any

import torch

from src.data.datasets import build_cheng_datasets
from src.eval.metrics import aggregate_scores, segmentation_scores
from src.eval.report import results_to_markdown, write_report
from src.models.registry import build_model, count_parameters

logger = logging.getLogger("evaluate")

#: Metrics where a lower value is better. Everything else is treated as
#: higher-is-better by the regression gate.
LOWER_IS_BETTER = {"hd95_mean", "hd95_std", "hd95_undefined", "inference_ms_per_case"}

#: Metrics the gate actually checks. Deliberately a short list of headline
#: numbers — gating on every column would make the build fail on noise.
GATED_METRICS = ("dice_mean", "iou_mean")


def load_checkpoint_model(
    checkpoint_path: Path, device: torch.device
) -> tuple[torch.nn.Module, dict]:
    """Rebuild a model from a checkpoint written by :class:`src.train.Trainer`.

    The training config travels inside the checkpoint, so the architecture is
    reconstructed from the checkpoint alone without needing the original YAML.
    """
    checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=False)
    if "config" not in checkpoint or "model_state" not in checkpoint:
        raise ValueError(
            f"{checkpoint_path} is not a NeuroSeg checkpoint "
            f"(expected 'config' and 'model_state' keys)"
        )

    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


@torch.no_grad()
def evaluate_segmentation_2d(
    model: torch.nn.Module,
    dataset,
    device: torch.device,
    batch_size: int = 16,
    threshold: float = 0.5,
) -> dict[str, float]:
    """Score a 2-D segmentation model over a dataset, per case."""
    from torch.utils.data import DataLoader

    loader = DataLoader(dataset, batch_size=batch_size, shuffle=False)
    bundles = []
    total_time = 0.0
    n_cases = 0

    for batch in loader:
        images = batch["image"].to(device)
        masks = batch["mask"].numpy() > 0.5

        start = time.perf_counter()
        logits = model(images)
        if device.type == "cuda":
            torch.cuda.synchronize()
        total_time += time.perf_counter() - start

        preds = (torch.sigmoid(logits.float()) > threshold).cpu().numpy()
        for pred, target in zip(preds, masks, strict=True):
            bundles.append(segmentation_scores(pred[0], target[0]))
            n_cases += 1

    results = aggregate_scores(bundles)
    results["inference_ms_per_case"] = 1000.0 * total_time / max(n_cases, 1)
    return results


#: Architectures the 2-D benchmark can evaluate. ``attention_unet`` is
#: dimension-parameterised, so its ``dim`` decides.
TWO_D_ARCHITECTURES = {"unet2d", "attention_unet"}


def is_2d_model(config: dict[str, Any]) -> bool:
    """Whether a checkpoint's architecture consumes ``(N, C, H, W)`` input."""
    model_config = config.get("model", {})
    name = str(model_config.get("name", "")).lower()
    if name not in TWO_D_ARCHITECTURES:
        return False
    return int(model_config.get("dim", 2)) == 2


def discover_checkpoints(artifacts_dir: Path, checkpoint_name: str = "best.pt") -> list[Path]:
    """Find every trained checkpoint under ``artifacts_dir``."""
    if not artifacts_dir.exists():
        return []
    return sorted(artifacts_dir.glob(f"*/{checkpoint_name}"))


def check_regressions(
    current: list[dict[str, Any]],
    baseline_path: Path,
    tolerance: float,
    metrics: tuple[str, ...] = GATED_METRICS,
) -> list[str]:
    """Compare results against a stored baseline.

    Returns:
        A list of human-readable failure messages; empty means the gate passes.
        A model present in the baseline but missing from the current run is
        itself a failure — silently dropping a model is exactly the kind of
        regression a gate should catch.
    """
    if not baseline_path.exists():
        logger.warning("no baseline at %s; skipping regression gate", baseline_path)
        return []

    baseline_rows = json.loads(baseline_path.read_text())
    baseline = {row["model"]: row for row in baseline_rows}
    current_by_name = {row["model"]: row for row in current}

    failures: list[str] = []

    for name, baseline_row in baseline.items():
        if name not in current_by_name:
            failures.append(f"{name}: present in baseline but missing from this run")
            continue

        row = current_by_name[name]
        for metric in metrics:
            if metric not in baseline_row or metric not in row:
                continue
            old, new = baseline_row[metric], row[metric]
            if not isinstance(old, (int, float)) or not isinstance(new, (int, float)):
                continue
            if old != old or new != new:  # NaN
                continue

            if metric in LOWER_IS_BETTER:
                delta = new - old
                if delta > tolerance:
                    failures.append(
                        f"{name}: {metric} rose {old:.4f} -> {new:.4f} (+{delta:.4f} > {tolerance})"
                    )
            else:
                delta = old - new
                if delta > tolerance:
                    failures.append(
                        f"{name}: {metric} fell {old:.4f} -> {new:.4f} (-{delta:.4f} > {tolerance})"
                    )

    return failures


def collect_auxiliary_results(artifacts_dir: Path) -> list[dict[str, Any]]:
    """Gather results written by the detection, 3-D and interpolation trainers.

    Those tasks are trained by their own scripts (they need different data and
    very different evaluation), and each writes a JSON summary next to its
    checkpoint. Rather than re-running them — which would mean reloading BraTS
    volumes and re-invoking Ultralytics — this reads what they recorded, so a
    single report can cover every model family.

    Missing files are skipped silently: a family that has not been trained
    simply does not appear in the report, which is the honest representation.
    """
    rows: list[dict[str, Any]] = []

    # 3-D segmentation (scripts/train_brats3d.py)
    for results_path in sorted(artifacts_dir.glob("*/results.json")):
        try:
            data = json.loads(results_path.read_text())
        except (OSError, json.JSONDecodeError):
            continue

        if "test" in data and isinstance(data["test"], dict):
            test = data["test"]
            rows.append(
                {
                    "model": results_path.parent.name,
                    "architecture": data.get("model", "?"),
                    "task": "segmentation-3d",
                    "params_M": round(data.get("params", 0) / 1e6, 2),
                    "dice_mean": test.get("dice_mean"),
                    "dice_whole": test.get("dice_whole"),
                    "dice_core": test.get("dice_core"),
                    "dice_enhancing": test.get("dice_enhancing"),
                    "n_cases": data.get("cases_per_split", {}).get("test"),
                }
            )
        elif "learned" in data:  # interpolation (scripts/train_interpolation.py)
            learned, baseline = data["learned"], data.get("linear_baseline", {})
            rows.append(
                {
                    "model": results_path.parent.name,
                    "architecture": "slice_interp_cnn",
                    "task": "interpolation",
                    "psnr_db": learned.get("psnr"),
                    "ssim": learned.get("ssim"),
                    "psnr_baseline_db": baseline.get("psnr"),
                    "ssim_baseline": baseline.get("ssim"),
                    "psnr_gain_db": data.get("psnr_gain_db"),
                    "n_cases": learned.get("n_slices"),
                }
            )

    # Detection (scripts/train_yolo.py)
    yolo_metrics = artifacts_dir / "yolo" / "test_metrics.json"
    if yolo_metrics.exists():
        try:
            data = json.loads(yolo_metrics.read_text())
            rows.append(
                {
                    "model": "yolov8_detection",
                    "architecture": "yolov8n",
                    "task": "detection",
                    "mAP50": data.get("mAP50"),
                    "mAP50_95": data.get("mAP50_95"),
                    "precision": data.get("precision"),
                    "recall": data.get("recall"),
                }
            )
        except (OSError, json.JSONDecodeError):
            pass

    return rows


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description="Benchmark NeuroSeg models")
    parser.add_argument(
        "--task",
        default="segmentation-2d",
        choices=["segmentation-2d", "all"],
        help="'segmentation-2d' evaluates 2-D checkpoints directly; 'all' adds "
        "the detection, 3-D and interpolation results recorded by their own "
        "training scripts",
    )
    parser.add_argument("--artifacts", type=Path, default=Path("artifacts"))
    parser.add_argument("--cache", type=Path, default=Path("data/processed/cheng_256.npz"))
    parser.add_argument("--split", default="test", choices=["train", "val", "test"])
    parser.add_argument("--out-dir", type=Path, default=Path("reports"))
    parser.add_argument("--name", default="benchmark")
    parser.add_argument("--subset", type=int, default=None, help="cap cases, for CI")
    parser.add_argument("--batch-size", type=int, default=16)
    parser.add_argument("--device", default="cuda")
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--gate", type=Path, default=None, help="baseline JSON to compare against")
    parser.add_argument("--tolerance", type=float, default=0.02)
    parser.add_argument("--save-baseline", type=Path, default=None)
    args = parser.parse_args(argv)

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

    device = torch.device(args.device if torch.cuda.is_available() else "cpu")
    if device.type == "cpu" and args.device == "cuda":
        logger.warning("CUDA unavailable; running on CPU")

    if not args.cache.exists():
        logger.error(
            "dataset cache not found at %s — run scripts/prepare_cheng.py first", args.cache
        )
        return 1

    checkpoints = discover_checkpoints(args.artifacts)
    if not checkpoints:
        logger.error("no checkpoints found under %s", args.artifacts)
        return 1

    datasets, splits = build_cheng_datasets(
        args.cache, seed=args.seed, augment_train=False, subset=args.subset
    )
    dataset = datasets[args.split]
    logger.info("evaluating on the %s split: %d slices", args.split, len(dataset))

    rows: list[dict[str, Any]] = []
    for checkpoint_path in checkpoints:
        name = checkpoint_path.parent.name
        try:
            model, config = load_checkpoint_model(checkpoint_path, device)
        except Exception as exc:  # noqa: BLE001
            logger.warning("skipping %s: %s", name, exc)
            continue

        # artifacts/ holds volumetric checkpoints too; feeding a 3-D model 2-D
        # Cheng slices raises deep inside the forward pass. Filter on the
        # architecture rather than letting it fail.
        if not is_2d_model(config):
            logger.info(
                "skipping %s: %s is volumetric, not a 2-D model "
                "(its results are collected via --task all)",
                name,
                config["model"].get("name", "?"),
            )
            del model
            continue

        logger.info("evaluating %s", name)
        results = evaluate_segmentation_2d(
            model,
            dataset,
            device,
            batch_size=args.batch_size,
            threshold=config.get("threshold", 0.5),
        )

        rows.append(
            {
                "model": name,
                "architecture": config["model"].get("name", "?"),
                "task": "segmentation-2d",
                "params_M": round(count_parameters(model) / 1e6, 2),
                **dict(results.items()),
            }
        )
        del model
        if device.type == "cuda":
            torch.cuda.empty_cache()

    if not rows:
        logger.error("no checkpoints could be evaluated")
        return 1

    rows.sort(key=lambda r: r.get("dice_mean", 0.0), reverse=True)

    auxiliary: list[dict[str, Any]] = []
    if args.task == "all":
        auxiliary = collect_auxiliary_results(args.artifacts)
        logger.info("collected %d result(s) from other model families", len(auxiliary))

    columns = [
        "model",
        "architecture",
        "task",
        "params_M",
        "dice_mean",
        "dice_std",
        "iou_mean",
        "hd95_mean",
        "hd95_undefined",
        "sensitivity_mean",
        "specificity_mean",
        "inference_ms_per_case",
        "n_cases",
    ]
    notes = [
        f"Split: Cheng {args.split}, {len(dataset)} slices, patient-disjoint (seed {args.seed}).",
        f"Device: {torch.cuda.get_device_name(0) if device.type == 'cuda' else 'CPU'}.",
        "Hausdorff distance is the 95th percentile, in pixels (the Cheng dataset "
        "ships no voxel spacing, so distances are not in mm).",
        "`hd95_undefined` counts cases where the model predicted an empty mask "
        "against a non-empty reference, making the distance undefined.",
    ]
    if args.subset:
        notes.append(
            f"**Subset run**: capped at {args.subset} slices per split — indicative only, "
            f"not a final result."
        )

    paths = write_report(rows, args.out_dir, args.name, columns, "NeuroSeg benchmark", notes)
    logger.info("report written to %s", paths["markdown"])

    if auxiliary:
        # Written as a separate table rather than merged into the one above:
        # detection mAP, reconstruction PSNR and segmentation Dice are not
        # comparable quantities, and putting them in one table with mostly-empty
        # cells invites exactly the cross-task comparison that would be wrong.
        by_task: dict[str, list[dict[str, Any]]] = {}
        for row in auxiliary:
            by_task.setdefault(row["task"], []).append(row)

        sections = [paths["markdown"].read_text(encoding="utf-8")]
        for task_name, task_rows in sorted(by_task.items()):
            task_columns = [
                key for key in task_rows[0] if any(r.get(key) is not None for r in task_rows)
            ]
            sections.append(
                results_to_markdown(task_rows, task_columns, f"{task_name} results", ())
            )

        combined = "\n\n".join(sections)
        paths["markdown"].write_text(combined, encoding="utf-8")
        (args.out_dir / f"{args.name}_all.json").write_text(
            json.dumps(rows + auxiliary, indent=2), encoding="utf-8"
        )

    print("\n" + paths["markdown"].read_text(encoding="utf-8"))

    if args.save_baseline:
        args.save_baseline.parent.mkdir(parents=True, exist_ok=True)
        args.save_baseline.write_text(json.dumps(rows, indent=2))
        logger.info("baseline saved to %s", args.save_baseline)

    if args.gate:
        failures = check_regressions(rows, args.gate, args.tolerance)
        if failures:
            logger.error("REGRESSION GATE FAILED:")
            for failure in failures:
                logger.error("  %s", failure)
            return 1
        logger.info("regression gate passed (tolerance %.3f)", args.tolerance)

    return 0


if __name__ == "__main__":
    sys.exit(main())