object_detection/evaluate.py

"""
evaluate.py -- Part B, Section 3.4 (evaluation and error analysis).

Reports on the frozen test split:
    mAP@0.50, mAP@0.75, mAP@0.95, mAP@[.50:.05:.95],
    precision / recall / F1 at a tuned confidence threshold,
    mean IoU of matched detections, inference time per image,
    recall stratified by polyp size (small / medium / large),
    a confidence-threshold sweep and a precision-recall curve,
    >= 8 seed-selected qualitative panels covering TP, FP and FN,
    and a failure-pattern summary linked back to the EDA.

Usage
-----
python evaluate.py --ckpt ../Results/model_weights/det_fasterrcnn_resnet50_fpn_ce_seed16.pt
python evaluate.py --compare det_fasterrcnn_resnet50_fpn_ce_seed16 \
                             det_fasterrcnn_resnet50_fpn_focal-a0.25g2.0_seed16
"""

from __future__ import annotations

import argparse
import json
import time
from pathlib import Path

import numpy as np
import pandas as pd
import torch

from common import (DET_SPLIT_CSV, METRICS_DIR, OVERLAY_DIR, PLOTS_DIR, SEED,
                    SEED_TAG, KvasirDetectionDataset, _mpl, average_precision,
                    box_iou_np, coco_map, get_device, make_loader,
                    prf_at_threshold, progress, save_json, seeded_rng, set_seed,
                    size_stratified_recall)
from model import build_detector, set_roi_loss


# ---------------------------------------------------------------------------
def _torch_load(path):
    """torch>=2.6 defaults weights_only=True, which cannot restore the config
    dict we store alongside the weights; older torch has no such argument."""
    try:
        return torch.load(path, map_location="cpu", weights_only=False)
    except TypeError:
        return torch.load(path, map_location="cpu")


def load_detector(ckpt_path: Path):
    ck = _torch_load(ckpt_path)
    cfg = ck["args"]
    set_roi_loss(cfg.get("loss", "ce"), alpha=cfg.get("focal_alpha", 0.25),
                 gamma=cfg.get("focal_gamma", 2.0), fg_weight=cfg.get("fg_weight", 4.0))
    model = build_detector(cfg.get("arch", "fasterrcnn_resnet50_fpn"),
                           pretrained_backbone=False,
                           img_size=cfg.get("img_size", 512),
                           trainable_backbone_layers=cfg.get("trainable_layers", 3))
    model.load_state_dict(ck["model_state"])
    model.eval()
    return model, ck, cfg


@torch.no_grad()
def infer_split(model, cfg, split: str, bs: int, workers: int):
    device = get_device()
    model.to(device).eval()
    ds = KvasirDetectionDataset(DET_SPLIT_CSV, split, augment=False,
                                img_size=cfg.get("img_size", 512))
    dl = make_loader(ds, bs, shuffle=False, num_workers=workers)

    preds, gts, stems, areas = [], [], [], []
    # warm-up so the timing excludes lazy CUDA/kernel initialisation
    for images, _ in dl:
        model([images[0].to(device)])
        break
    if device.type == "cuda":
        torch.cuda.synchronize()

    t0 = time.time()
    for images, targets in progress(dl, desc=f"infer[{split}]"):
        outs = model([im.to(device) for im in images])
        for im, o, t in zip(images, outs, targets):
            preds.append({"boxes": o["boxes"].cpu().numpy(),
                          "scores": o["scores"].cpu().numpy()})
            gts.append({"boxes": t["boxes"].numpy()})
            stems.append(t["stem"])
            areas.append(float(im.shape[-1] * im.shape[-2]))
    if device.type == "cuda":
        torch.cuda.synchronize()
    dt = time.time() - t0
    timing = {"total_seconds": dt, "n_images": len(stems),
              "ms_per_image": 1000 * dt / max(1, len(stems)),
              "fps": len(stems) / max(1e-9, dt), "device": str(device),
              "batch_size": bs}
    return preds, gts, stems, areas, timing


# ---------------------------------------------------------------------------
def threshold_sweep(preds, gts, iou_thr: float = 0.5,
                    thresholds=np.arange(0.05, 0.96, 0.05)) -> pd.DataFrame:
    rows = [prf_at_threshold(preds, gts, iou_thr, float(t)) for t in thresholds]
    return pd.DataFrame(rows)


def plot_pr_and_sweep(preds, gts, sweep: pd.DataFrame, tag: str):
    plt = _mpl()
    ap = average_precision(preds, gts, 0.5)
    fig, axes = plt.subplots(1, 2, figsize=(11.5, 4.4))
    if "recall_curve" in ap:
        axes[0].plot(ap["recall_curve"], ap["precision_curve"], lw=2, color="#4C72B0")
    axes[0].set_xlabel("recall"); axes[0].set_ylabel("precision")
    axes[0].set_title(f"PR curve @IoU 0.50 (AP={ap['ap']:.3f})")
    axes[0].set_xlim(0, 1); axes[0].set_ylim(0, 1.02); axes[0].grid(alpha=0.3)

    axes[1].plot(sweep["score_threshold"], sweep["precision"], marker="o", ms=3, label="precision")
    axes[1].plot(sweep["score_threshold"], sweep["recall"], marker="s", ms=3, label="recall")
    axes[1].plot(sweep["score_threshold"], sweep["f1"], marker="^", ms=3, label="F1")
    best = sweep.loc[sweep["f1"].idxmax()]
    axes[1].axvline(best["score_threshold"], ls="--", c="k", lw=1)
    axes[1].set_xlabel("confidence threshold")
    axes[1].set_title(f"Operating point sweep (best F1={best['f1']:.3f} "
                      f"@ {best['score_threshold']:.2f})")
    axes[1].legend(); axes[1].grid(alpha=0.3)
    fig.suptitle(tag, fontsize=11)
    fig.tight_layout()
    out = PLOTS_DIR / f"pr_sweep_{tag}.png"
    fig.savefig(out, dpi=160); plt.close(fig)
    print(f"[saved] {out}")


def qualitative(preds, gts, stems, cfg, tag: str, score_thr: float,
                iou_thr: float = 0.5, n: int = 8):
    """Seed-selected panels guaranteed to contain true positives, false
    positives and false negatives."""
    plt = _mpl()
    from PIL import Image

    df = pd.read_csv(DET_SPLIT_CSV).set_index("image_id")
    per = []
    for i, (p, g, s) in enumerate(zip(preds, gts, stems)):
        keep = p["scores"] >= score_thr
        pb, ps = p["boxes"][keep], p["scores"][keep]
        gb = g["boxes"]
        ious = box_iou_np(pb, gb) if len(pb) and len(gb) else np.zeros((len(pb), len(gb)))
        tp = int((ious.max(axis=1) >= iou_thr).sum()) if ious.size else 0
        fp = len(pb) - tp
        fn = len(gb) - (int((ious.max(axis=0) >= iou_thr).sum()) if ious.size else 0)
        per.append({"i": i, "stem": s, "tp": tp, "fp": fp, "fn": fn,
                    "n_pred": len(pb), "n_gt": len(gb),
                    "best_iou": float(ious.max()) if ious.size else 0.0})
    pdf = pd.DataFrame(per)

    rng = seeded_rng(offset=9)
    picks: list[int] = []

    def take(sub, k):
        sub = sub[~sub["i"].isin(picks)]
        if len(sub) == 0:
            return
        idx = rng.choice(len(sub), size=min(k, len(sub)), replace=False)
        picks.extend(sub.iloc[idx]["i"].tolist())

    take(pdf[(pdf.tp > 0) & (pdf.fp == 0) & (pdf.fn == 0)], 3)   # clean successes
    take(pdf[pdf.fp > 0], 3)                                      # false positives
    take(pdf[pdf.fn > 0], 3)                                      # missed polyps
    take(pdf, n)
    picks = picks[:max(n, 8)]

    ncol = 4
    nrow = int(np.ceil(len(picks) / ncol))
    fig, axes = plt.subplots(nrow, ncol, figsize=(4.2 * ncol, 4.5 * nrow))
    axes = np.atleast_1d(axes).ravel()
    recs = []
    for k, i in enumerate(picks):
        stem = stems[i]
        r = df.loc[stem]
        img = Image.open(r["image_path"]).convert("RGB")
        size = cfg.get("img_size", 512)
        img = img.resize((size, size), Image.BILINEAR)
        axes[k].imshow(np.asarray(img))
        for b in gts[i]["boxes"]:
            axes[k].add_patch(plt.Rectangle((b[0], b[1]), b[2] - b[0], b[3] - b[1],
                                            fill=False, edgecolor="lime", lw=2.2))
        keep = preds[i]["scores"] >= score_thr
        pbs, pss = preds[i]["boxes"][keep], preds[i]["scores"][keep]
        ious = box_iou_np(pbs, gts[i]["boxes"]) if len(pbs) and len(gts[i]["boxes"]) \
            else np.zeros((len(pbs), max(1, len(gts[i]["boxes"]))))
        for j, (b, sc) in enumerate(zip(pbs, pss)):
            best = float(ious[j].max()) if ious.size else 0.0
            col = "red" if best >= iou_thr else "orange"
            axes[k].add_patch(plt.Rectangle((b[0], b[1]), b[2] - b[0], b[3] - b[1],
                                            fill=False, edgecolor=col, lw=1.8, ls="-"))
            axes[k].text(b[0], max(10, b[1] - 4), f"{sc:.2f} / IoU {best:.2f}",
                         color=col, fontsize=7,
                         bbox=dict(facecolor="black", alpha=0.45, pad=0.8, edgecolor="none"))
        row = pdf.iloc[i]
        axes[k].set_title(f"{stem[:18]}\nTP={row.tp} FP={row.fp} FN={row.fn}", fontsize=8)
        axes[k].axis("off")
        recs.append({"image_id": stem, "tp": int(row.tp), "fp": int(row.fp),
                     "fn": int(row.fn), "n_pred": int(row.n_pred), "n_gt": int(row.n_gt),
                     "best_iou": round(float(row.best_iou), 3),
                     "max_confidence": round(float(pss.max()), 3) if len(pss) else 0.0})
    for k in range(len(picks), len(axes)):
        axes[k].axis("off")
    fig.suptitle(f"Kvasir-SEG test predictions -- {tag} "
                 f"(green = GT, red = matched detection, orange = false positive; "
                 f"conf>={score_thr:.2f}, seed={SEED})", fontsize=11)
    fig.tight_layout()
    out = OVERLAY_DIR / f"det_examples_{tag}.png"
    fig.savefig(out, dpi=145); plt.close(fig)
    print(f"[saved] {out}")

    pd.DataFrame(recs).to_csv(METRICS_DIR / f"det_examples_{tag}.csv", index=False)
    return pdf


def failure_patterns(pdf: pd.DataFrame, preds, gts, areas, score_thr: float) -> dict:
    """Quantified failure taxonomy, each row linkable to an EDA finding."""
    n_img = len(pdf)
    small_missed = medium_missed = large_missed = 0
    for p, g, a in zip(preds, gts, areas):
        keep = p["scores"] >= score_thr
        pb = p["boxes"][keep]
        gb = g["boxes"]
        if len(gb) == 0:
            continue
        ious = box_iou_np(pb, gb) if len(pb) else np.zeros((0, len(gb)))
        best = ious.max(axis=0) if len(pb) else np.zeros(len(gb))
        for j in range(len(gb)):
            if best[j] >= 0.5:
                continue
            frac = (gb[j, 2] - gb[j, 0]) * (gb[j, 3] - gb[j, 1]) / max(1.0, a)
            if frac < 0.01:
                small_missed += 1
            elif frac < 0.05:
                medium_missed += 1
            else:
                large_missed += 1
    return {
        "images_evaluated": int(n_img),
        "images_fully_correct": int(((pdf.fp == 0) & (pdf.fn == 0) & (pdf.tp > 0)).sum()),
        "images_with_false_positive": int((pdf.fp > 0).sum()),
        "images_with_missed_polyp": int((pdf.fn > 0).sum()),
        "images_with_no_detection_at_all": int((pdf.n_pred == 0).sum()),
        "images_with_duplicate_detections": int((pdf.n_pred > pdf.n_gt).sum()),
        "missed_by_size": {"small_lt_1pct": small_missed,
                           "medium_1_to_5pct": medium_missed,
                           "large_ge_5pct": large_missed},
        "localisation_near_misses_iou_0.3_to_0.5":
            int(((pdf.best_iou >= 0.3) & (pdf.best_iou < 0.5)).sum()),
    }


# ---------------------------------------------------------------------------
def evaluate_one(args):
    model, ck, cfg = load_detector(Path(args.ckpt))
    tag = ck.get("tag", Path(args.ckpt).stem)
    preds, gts, stems, areas, timing = infer_split(model, cfg, args.split, args.bs, args.workers)

    sweep = threshold_sweep(preds, gts, args.iou_thr)
    best_row = sweep.loc[sweep["f1"].idxmax()]
    score_thr = args.score_thr if args.score_thr is not None else float(best_row["score_threshold"])

    maps = coco_map(preds, gts)
    op = prf_at_threshold(preds, gts, args.iou_thr, score_thr)
    sizes = size_stratified_recall(preds, gts, areas, args.iou_thr, score_thr)

    out = {"tag": tag, "split": args.split, "seed": SEED,
           "loss_description": ck.get("loss_description", cfg.get("loss")),
           "confidence_threshold_used": score_thr,
           "confidence_threshold_selection": "argmax F1 on the sweep"
                                             if args.score_thr is None else "user supplied",
           "mAP@0.50": maps["AP@0.50"], "mAP@0.75": maps["AP@0.75"],
           "mAP@0.95": maps["AP@0.95"],
           "mAP@[.50:.05:.95]": maps["mAP@[.50:.05:.95]"],
           "AP_per_iou_threshold": maps,
           "operating_point": op,
           "recall_by_polyp_size": sizes,
           "timing": timing, "checkpoint": str(args.ckpt), "config": cfg}

    pdf = qualitative(preds, gts, stems, cfg, tag, score_thr, args.iou_thr, args.n_examples)
    out["failure_patterns"] = failure_patterns(pdf, preds, gts, areas, score_thr)

    save_json(out, METRICS_DIR / f"eval_{tag}_{args.split}.json")
    sweep.to_csv(METRICS_DIR / f"sweep_{tag}_{args.split}.csv", index=False)
    pdf.to_csv(METRICS_DIR / f"per_image_{tag}_{args.split}.csv", index=False)
    plot_pr_and_sweep(preds, gts, sweep, tag)

    print(f"\n=== {tag} / {args.split} ===")
    print(f"  mAP@0.50            {out['mAP@0.50']:.4f}")
    print(f"  mAP@0.75            {out['mAP@0.75']:.4f}")
    print(f"  mAP@0.95            {out['mAP@0.95']:.4f}")
    print(f"  mAP@[.50:.05:.95]   {out['mAP@[.50:.05:.95]']:.4f}")
    print(f"  @conf={score_thr:.2f}: P={op['precision']:.4f} R={op['recall']:.4f} "
          f"F1={op['f1']:.4f}")
    print(f"  TP={op['true_positives']} FP={op['false_positives']} FN={op['false_negatives']}")
    print(f"  mean IoU (matched)  {op['mean_iou_matched']:.4f}")
    print(f"  inference           {timing['ms_per_image']:.1f} ms/image "
          f"({timing['fps']:.1f} FPS on {timing['device']})")
    print("  recall by polyp size:")
    for k, v in sizes.items():
        print(f"    {k:8s} n={v['n_gt']:4d}  recall={v['recall']:.4f}")
    print("  failure patterns:", json.dumps(out["failure_patterns"], indent=4))


def compare(tags, split: str = "test"):
    """Section 3.3.4: how the loss-function ablation moves each metric."""
    rows = []
    for t in tags:
        p = METRICS_DIR / f"eval_{t}_{split}.json"
        if not p.exists():
            print(f"[warn] missing {p}")
            continue
        m = json.loads(p.read_text())
        op, sz = m["operating_point"], m["recall_by_polyp_size"]
        rows.append({
            "run": t, "loss": m.get("loss_description", "")[:46],
            "mAP@0.50": m["mAP@0.50"], "mAP@0.75": m["mAP@0.75"],
            "mAP@[.50:.95]": m["mAP@[.50:.05:.95]"],
            "precision": op["precision"], "recall": op["recall"], "F1": op["f1"],
            "FP": op["false_positives"], "FN": op["false_negatives"],
            "mean_IoU": op["mean_iou_matched"],
            "recall_small": sz["small"]["recall"], "recall_medium": sz["medium"]["recall"],
            "recall_large": sz["large"]["recall"],
            "ms_per_image": m["timing"]["ms_per_image"]})
    if not rows:
        return
    df = pd.DataFrame(rows)
    out = METRICS_DIR / f"ablation_comparison_{split}_{SEED_TAG}.csv"
    df.round(4).to_csv(out, index=False)
    print(f"\n=== Loss-function ablation ({split}) ===")
    print(df.round(4).to_string(index=False))
    if len(df) >= 2:
        base, alt = df.iloc[0], df.iloc[1]
        print(f"\nDelta ({alt['run']} - {base['run']}):")
        for k in ["mAP@0.50", "mAP@[.50:.95]", "precision", "recall", "F1",
                  "recall_small", "recall_medium", "recall_large"]:
            print(f"  {k:16s} {alt[k] - base[k]:+.4f}")
        print(f"  {'FP':16s} {alt['FP'] - base['FP']:+d}")
        print(f"  {'FN':16s} {alt['FN'] - base['FN']:+d}")
    print(f"[saved] {out}")


def main():
    ap = argparse.ArgumentParser(description="Kvasir-SEG detector evaluation (seed 16)")
    ap.add_argument("--ckpt")
    ap.add_argument("--split", default="test", choices=["train", "val", "test"])
    ap.add_argument("--bs", type=int, default=4)
    ap.add_argument("--workers", type=int, default=2)
    ap.add_argument("--iou-thr", type=float, default=0.5)
    ap.add_argument("--score-thr", type=float, default=None,
                    help="fixed confidence threshold; default picks argmax-F1 from the sweep")
    ap.add_argument("--n-examples", type=int, default=8)
    ap.add_argument("--compare", nargs="*", default=None)
    args = ap.parse_args()

    set_seed(SEED)
    if args.compare:
        compare(args.compare, args.split)
        return
    if not args.ckpt:
        ap.error("--ckpt is required (or use --compare)")
    evaluate_one(args)


if __name__ == "__main__":
    main()