AareLC_train/tools/compare_unseen.py

#!/usr/bin/env python3
"""Every measured model on the gold frames none of them trained on, with the noise made visible.

    python tools/compare_unseen.py            # table, runs/clean_gold_comparison.json

Two sets of frames:

* the **core**: the 149 frames that were unseen by every model on 2026-10-02
  (``data/gold/unseen_core.json``). It never changes, so the progression of the models over time
  stays comparable;
* **all unseen** frames: the core plus every gold frame added since that no dataset has used.
  It grows with the gold set and is the one to judge a candidate on.

A model is listed when its predictions on the current gold set exist under
``runs/_measure_<version>/gold_predictions`` (``train_run.py`` keeps them for a new run; after
the gold set grew, re-measure the older versions or they are compared on fewer frames).
For every model the script also reports, against the active model, the paired-bootstrap
interval of the difference in F1 and mask score and the share of draws it wins (tools/gate_stats.py).
"""
from __future__ import annotations

import argparse
import json
import re
import sys
from pathlib import Path
from typing import Dict, List, Optional

HERE = Path(__file__).resolve().parent
sys.path.insert(0, str(HERE))
import gate_stats  # noqa: E402

NAMES = ["loop_all", "pin", "crystal", "loop_face", "ice", "needle"]
BEAMLINE_PREFIX = "yolo26n-seg-overlap-false"   # the model trained before the loop


def class_report(c) -> Dict[str, float]:
    tp, fp, fn, iou = c
    precision = tp / (tp + fp) if (tp + fp) else 0.0
    recall = tp / (tp + fn) if (tp + fn) else 0.0
    f1 = 2 * precision * recall / (precision + recall) if (precision + recall) else 0.0
    return {"tp": int(tp), "fp": int(fp), "fn": int(fn), "precision": precision, "recall": recall, "f1": f1,
            "mean_mask_iou": iou / tp if tp else 0.0, "n_matched": int(tp)}


def report(frames: Dict[str, gate_stats.FrameCounts]) -> Dict:
    rows = list(frames.values())
    per_class = {}
    for cid, name in enumerate(NAMES):
        c = gate_stats._total(rows, cid)
        if c[0] or c[1] or c[2]:
            per_class[name] = class_report(c)
    return {"per_class": per_class, "overall": class_report(gate_stats._total(rows, None)),
            "dataset": {"gold_images": len(rows)}}


def mask_score(o: Dict) -> float:
    return o["mean_mask_iou"] * (o["tp"] / (o["tp"] + o["fn"]) if (o["tp"] + o["fn"]) else 0.0)


def table(title: str, reps: Dict[str, Dict]) -> None:
    print(title)
    print(f"{'run':46} {'frames':>6} {'F1':>6} {'recall':>7} {'IoU matched':>12} {'mask score':>11}   loop/pin/crystal/face/ice F1")
    for k, r in reps.items():
        o, pc = r["overall"], r["per_class"]
        print(f"{k:46} {r['dataset']['gold_images']:6d} {o['f1']:6.3f} {o['recall']:7.2f} {o['mean_mask_iou']:12.3f} {mask_score(o):11.3f}   "
              + "/".join(f"{pc.get(c, {}).get('f1', 0):.2f}" for c in NAMES[:5]))


def main(argv: Optional[List[str]] = None) -> int:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--root", default=str(HERE.parent), help="AareLC_train checkout (default: this one)")
    ap.add_argument("--draws", type=int, default=2000)
    args = ap.parse_args(argv)
    root = Path(args.root).expanduser()
    sys.path.insert(0, str(root))
    gold = root / "data" / "gold" / "labels"
    core_file = root / "data" / "gold" / "unseen_core.json"
    if not core_file.exists():
        print(f"{core_file} is missing: it lists the frozen core and the gold frames an older model trained on")
        return 2
    cj = json.loads(core_file.read_text(encoding="utf-8"))
    core, seen = set(cj["core"]), set(cj.get("seen", []))
    trained_on = set()
    for m in (root / "data" / "datasets").glob("*/manifest.csv"):
        with open(m, encoding="utf-8") as fh:
            next(fh, None)
            for line in fh:
                key = line.split(",")[0].strip()
                if key.startswith("tpi_"):
                    trained_on.add(key)
    stems = sorted(p.stem for p in gold.glob("*.txt"))
    key_of = lambda s: "tpi_" + re.search(r"tpi_(\d+)", s).group(1)
    unseen = [s for s in stems if s not in seen and key_of(s) not in trained_on]
    core_now = [s for s in unseen if s in core]

    ver2run = {}
    for rj in (root / "runs").glob("*/run.json"):
        r = json.loads(rj.read_text(encoding="utf-8"))
        if r.get("version"):
            ver2run[r["version"]] = r["id"]
    active = None
    ap_file = root / "models" / "active" / "production.json"
    if ap_file.exists():
        active = json.loads(ap_file.read_text(encoding="utf-8")).get("version")

    frames: Dict[str, Dict[str, gate_stats.FrameCounts]] = {}
    label_of: Dict[str, str] = {}
    for d in sorted((root / "runs").glob("_measure_*/gold_predictions")):
        ver = d.parent.name[len("_measure_"):]
        label = ("beamline " + ver) if ver.startswith(BEAMLINE_PREFIX) else ver2run.get(ver, ver)
        frames[label] = gate_stats.per_frame(d, gold, stems=unseen)
        label_of[ver] = label
    order = sorted(frames, key=lambda k: (not k.startswith("beamline"), k))
    reps = {k: report(frames[k]) for k in order}
    core_reps = {k: report({s: f for s, f in frames[k].items() if s in core}) for k in order}

    print(f"clean gold frames: {len(unseen)}" + (f" (core {len(core_now)})" if len(unseen) != len(core_now) else ""))
    table("", reps)
    short = [k for k in order if reps[k]["dataset"]["gold_images"] < len(unseen)]
    if short:
        print(f"!! measured on fewer frames than the current unseen set (re-measure them): {', '.join(short)}")
    if len(unseen) != len(core_now):
        table(f"\ncore, the {len(core_now)} frames frozen on {cj.get('frozen', '?')}:", core_reps)

    boot = {}
    base_label = label_of.get(active) if active else None
    if base_label and base_label in frames:
        print(f"\nagainst the active model ({base_label}), paired bootstrap over the frames both have: difference [95 % interval], share of draws won")
        classes = {"crystal": NAMES.index("crystal")}
        for k in order:
            if k == base_label:
                continue
            b = gate_stats.paired_bootstrap(frames[k], frames[base_label], draws=args.draws, classes=classes)
            if not b.get("frames"):
                continue
            boot[k] = b
            f, m, c = b["f1"], b["mask_score"], b["per_class"]["crystal"]["f1"]
            print(f"{k:46} F1 {f['delta']:+.3f} [{f['lo']:+.3f}, {f['hi']:+.3f}] {f['win']:4.0%}   mask {m['delta']:+.3f} [{m['lo']:+.3f}, {m['hi']:+.3f}] {m['win']:4.0%}   crystal {c['delta']:+.3f} [{c['lo']:+.3f}, {c['hi']:+.3f}] {c['win']:4.0%}")

    # per frame and model: overall and crystal counts, so any two models can be compared head to head
    # with their noise interval (the monitor page runs the paired bootstrap in the browser)
    crystal = NAMES.index("crystal")

    def row(f):
        o = gate_stats._total([f], None); c = gate_stats._total([f], crystal)
        return [int(o[0]), int(o[1]), int(o[2]), round(o[3], 3), int(c[0]), int(c[1]), int(c[2]), round(c[3], 3)]
    per_frame = {"stems": unseen, "fields": ["tp", "fp", "fn", "iou_sum", "crystal_tp", "crystal_fp", "crystal_fn", "crystal_iou_sum"],
                 "models": {k: [row(frames[k][st]) if st in frames[k] else None for st in unseen] for k in order}}
    (root / "runs" / "_unseen_per_frame.json").write_text(json.dumps(per_frame, separators=(",", ":")), encoding="utf-8")

    out = {"clean_frames": len(unseen), "reports": reps, "core_frames": len(core_now), "core_frozen": cj.get("frozen"),
           "core_reports": core_reps, "active": active, "bootstrap_vs_active": boot}
    (root / "runs" / "clean_gold_comparison.json").write_text(json.dumps(out, indent=2), encoding="utf-8")
    return 0


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