AareLC_train/tools/gate_stats.py

"""How sure is a difference between two models measured on the same gold frames?

The gold set is small (149 unseen frames on 2026-10-02: 40 crystals, 10 ice objects), so a
difference of one or two points in F1 between two models can be the luck of which frames are
gold. A paired bootstrap answers that: draw the gold frames again with replacement, the SAME
draw for both models, and look at the spread of the difference. ``win`` is the share of draws
in which the candidate beats the baseline.

Frames of one mount are not independent (the same loop at another angle, 24 of the 203 gold
frames on 2026-10-02), so the draw is by mount when ``data/gold/groups.json`` names them: a
mount is drawn with all its frames. Frame-wise draws would call the interval narrower than it is.

Pure Python on purpose (the gate's tests run without numpy); only ``per_frame`` needs the
evaluation code, and imports it when called.
"""
from __future__ import annotations

import json
import random
from pathlib import Path
from typing import Dict, Iterable, List, Mapping, Optional, Sequence, Tuple

Counts = Tuple[float, float, float, float]          # tp, fp, fn, sum of matched mask IoU
FrameCounts = Dict[int, Counts]                      # class id -> counts, for one frame


def load_groups(gold_root) -> Dict[str, Dict[str, str]]:
    """``data/gold/groups.json`` (tools/gold_groups.py): ``mount`` = frame -> mount id (a frame not
    listed is its own mount), ``light`` = frame -> state for the frames whose light is not usable,
    ``excluded`` = frame -> why it is left out of every measurement. Empty when the file is missing."""
    p = Path(gold_root) / "groups.json"
    g = json.loads(p.read_text(encoding="utf-8")) if p.exists() else {}
    return {"mount": dict(g.get("mount") or {}), "light": dict(g.get("light") or {}), "excluded": dict(g.get("excluded") or {})}


def per_frame(pred_dir, gold_dir, iou_thr: float = 0.5, stems: Optional[Iterable[str]] = None) -> Dict[str, FrameCounts]:
    """Per gold frame and class: tp, fp, fn and the summed mask IoU of the matches.

    ``stems`` restricts the frames; a gold frame without a prediction file is left out
    (the caller pairs on the frames both models have)."""
    from src.eval.metrics import evaluate_image
    from src.eval.yolo_labels import read_label_file

    pred_dir, gold_dir = Path(pred_dir), Path(gold_dir)
    wanted = set(stems) if stems is not None else None
    out: Dict[str, FrameCounts] = {}
    for g in sorted(gold_dir.glob("*.txt")):
        if wanted is not None and g.stem not in wanted:
            continue
        p = pred_dir / g.name
        if not p.exists():
            continue
        frame: FrameCounts = {}
        for c, cc in evaluate_image(read_label_file(g), read_label_file(p), iou_thr).items():
            frame[int(c)] = (float(cc.tp), float(cc.fp), float(cc.fn), float(sum(cc.mask_ious)))
        out[g.stem] = frame
    return out


def _total(frames: Sequence[FrameCounts], class_id: Optional[int]) -> Counts:
    tp = fp = fn = iou = 0.0
    for f in frames:
        for c, (a, b, d, e) in f.items():
            if class_id is None or c == class_id:
                tp += a; fp += b; fn += d; iou += e
    return tp, fp, fn, iou


def f1_of(c: Counts) -> float:
    tp, fp, fn, _ = c
    return 2 * tp / (2 * tp + fp + fn) if tp else 0.0


def mask_score_of(c: Counts) -> float:
    """Mean mask IoU over every gold object, a missed object counting as zero."""
    tp, _, fn, iou = c
    return iou / (tp + fn) if (tp + fn) else 0.0


def _summary(point: float, deltas: List[float]) -> Dict[str, float]:
    d = sorted(deltas)
    n = len(d)
    lo, hi = d[int(0.025 * (n - 1))], d[int(0.975 * (n - 1))]
    return {"delta": round(point, 4), "lo": round(lo, 4), "hi": round(hi, 4),
            "win": round(sum(1 for x in d if x > 0) / n, 3)}


def paired_bootstrap(cand: Mapping[str, FrameCounts], base: Mapping[str, FrameCounts], *, draws: int = 2000,
                     seed: int = 0, classes: Optional[Mapping[str, int]] = None,
                     groups: Optional[Mapping[str, str]] = None) -> Dict:
    """Candidate minus baseline on the frames both have: point difference, 95 % interval and win share
    for F1 and the mask score, overall and for each class in ``classes`` (name -> class id).

    ``groups`` (frame -> mount id) makes the draw by mount: the frames of a mount come and go together."""
    keys = sorted(set(cand) & set(base))
    n = len(keys)
    if n == 0:
        return {"frames": 0, "draws": 0}
    a = [cand[k] for k in keys]; b = [base[k] for k in keys]
    members: Dict[str, List[int]] = {}
    for i, k in enumerate(keys):
        members.setdefault((groups or {}).get(k) or f"frame:{k}", []).append(i)
    units = [members[g] for g in sorted(members)]
    targets: List[Tuple[str, Optional[int]]] = [("overall", None)] + [(name, cid) for name, cid in (classes or {}).items()]
    point = {}
    for name, cid in targets:
        ca, cb = _total(a, cid), _total(b, cid)
        point[name] = (f1_of(ca) - f1_of(cb), mask_score_of(ca) - mask_score_of(cb))
    rng = random.Random(seed)
    d_f1 = {name: [] for name, _ in targets}; d_ms = {name: [] for name, _ in targets}
    for _ in range(draws):
        idx = [i for _ in range(len(units)) for i in units[rng.randrange(len(units))]]
        sa = [a[i] for i in idx]; sb = [b[i] for i in idx]
        for name, cid in targets:
            ca, cb = _total(sa, cid), _total(sb, cid)
            d_f1[name].append(f1_of(ca) - f1_of(cb)); d_ms[name].append(mask_score_of(ca) - mask_score_of(cb))
    out = {"frames": n, "mounts": len(units), "draws": draws,
           "f1": _summary(point["overall"][0], d_f1["overall"]),
           "mask_score": _summary(point["overall"][1], d_ms["overall"])}
    if classes:
        out["per_class"] = {name: {"f1": _summary(point[name][0], d_f1[name]), "mask_score": _summary(point[name][1], d_ms[name])}
                            for name in classes}
    return out