scene-actor-extraction/scripts/validation/sample_eval.py

#!/usr/bin/env python3
"""
sample_eval.py — offline per-scene presence eval by timepoint sampling.

Annealing (anneal_sec) means an actor is "present" only after the whole file is
merged into [start,end] windows, so we cannot score live: we process → write the
pipeline JSON → sample timepoints → compare predicted vs ground-truth presence
sets → micro-sum TP/FP/FN → precision / recall / F1.
See [[per-scene-presence-eval-design]].

Usage:
  # against Amazon X-Ray CSVs (Zenodo)
  python scripts/validation/sample_eval.py \
      --pred "Scene in a Mall.json" \
      --xray  /data/xray/tt0384766 \
      --step 1.0

  # against MovieNet-PS (needs the .mat split + a title tt-id)
  python scripts/validation/sample_eval.py \
      --pred out.json \
      --movienet /data/movienet --split Train_app10 --title tt0032138

Sampling:
  --step S        regular grid every S seconds (default 1.0) — time-weighted headline
  --random N      N uniform-random timepoints instead of a grid (for CIs)
  --scene-anchored  one timepoint at each GT scene midpoint (X-Ray "per-scene" question)

Masking: scoring is restricted to actors present in BOTH the pipeline gallery
(--gallery) AND the ground truth. A GT actor absent from the gallery is ignored
(not counted as a miss) so we measure pipeline accuracy, not gallery coverage.
Pass --no-mask to disable.
"""

from __future__ import annotations

import argparse
import json
import random
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent))
from identity import keys_for                       # noqa: E402
from ground_truth import XRayGroundTruth, MovieNetGroundTruth  # noqa: E402
from tmdb_imdb_map import CrosswalkTable             # noqa: E402


# ── pipeline output → presence timeline ────────────────────────────────────────

class Prediction:
    """Pipeline output (minimal/standard schema) as per-actor presence windows."""

    def __init__(self, path: str | Path, crosswalk=None) -> None:
        with open(path) as f:
            data = json.load(f)
        self.movie = data.get("movie", "")
        self.anneal_sec = data.get("anneal_sec")
        self.actors: list[dict] = []
        self._max_t = 0.0
        for a in data.get("actors", []):
            keys = keys_for(imdb_id=a.get("imdb_id"), tmdb_id=a.get("tmdb_id"),
                            jellyfin_id=a.get("jellyfin_id"), name=a.get("name"),
                            crosswalk=crosswalk)
            windows = [(float(t0), float(t1)) for t0, t1 in a.get("scenes", [])]
            for _, t1 in windows:
                self._max_t = max(self._max_t, t1)
            self.actors.append({"keys": keys, "windows": windows})

    def present_at(self, t: float) -> set[frozenset[str]]:
        """Set of actors present at t; each actor is its (frozen) key-set."""
        out: set[frozenset[str]] = set()
        for a in self.actors:
            for t0, t1 in a["windows"]:
                if t0 <= t <= t1:
                    out.add(frozenset(a["keys"]))
                    break
        return out

    def present_in_span(self, s0: float, s1: float) -> set[frozenset[str]]:
        """Actors with ANY detection window overlapping [s0,s1].

        Snaps detections to a scene grid: an actor seen anywhere inside a scene
        counts as present for the whole scene. Isolates 'did we see this actor in
        this scene at all' (coverage) from exact-timing recall."""
        out: set[frozenset[str]] = set()
        for a in self.actors:
            for t0, t1 in a["windows"]:
                if t0 <= s1 and t1 >= s0:  # interval overlap
                    out.add(frozenset(a["keys"]))
                    break
        return out

    def all_keys(self) -> set[str]:
        out: set[str] = set()
        for a in self.actors:
            out |= a["keys"]
        return out

    @property
    def max_t(self) -> float:
        return self._max_t


def load_gallery_keys(path: str | None, crosswalk=None) -> set[str] | None:
    """Union of match keys for every actor in the gallery, for masking.

    Accepts either the JSON gallery or the HDF5 fast-load gallery (.h5/.hdf5,
    produced by json_to_hdf5_gallery.py) — the matcher reads HDF5, so this side must
    too. HDF5 stores ids/names as parallel string datasets."""
    if not path:
        return None
    out: set[str] = set()
    if path.endswith(".h5") or path.endswith(".hdf5"):
        import h5py
        with h5py.File(path, "r") as f:
            def col(name):
                return [(v.decode() if isinstance(v, bytes) else str(v))
                        for v in f[name][:]] if name in f else []
            imdb, tmdb = col("imdb_id"), col("tmdb_id")
            jf, name = col("jellyfin_id"), col("name")
            for i in range(len(name)):
                out |= keys_for(imdb_id=imdb[i] if i < len(imdb) else "",
                                tmdb_id=tmdb[i] if i < len(tmdb) else "",
                                jellyfin_id=jf[i] if i < len(jf) else "",
                                name=name[i], crosswalk=crosswalk)
        return out
    with open(path) as f:
        data = json.load(f)
    for a in data.get("actors", []):
        out |= keys_for(imdb_id=a.get("imdb_id"), tmdb_id=a.get("tmdb_id"),
                        jellyfin_id=a.get("jellyfin_id"), name=a.get("name"),
                        crosswalk=crosswalk)
    return out


# ── sampling ────────────────────────────────────────────────────────────────

def sample_points(args, pred: Prediction, gt) -> list[float]:
    if args.scene_anchored:
        spans = gt.scene_windows()
        if not spans:
            sys.exit("[eval] --scene-anchored: ground truth has no scene spans")
        return [(t0 + t1) / 2.0 for t0, t1 in spans]

    end = args.end if args.end is not None else max(pred.max_t, _gt_end(gt))
    if end <= 0:
        sys.exit("[eval] could not determine timeline end; pass --end")

    if args.random:
        rng = random.Random(args.seed)
        return sorted(rng.uniform(0.0, end) for _ in range(args.random))

    n = int(end / args.step) + 1
    return [i * args.step for i in range(n)]


def _gt_end(gt) -> float:
    spans = gt.scene_windows()
    return max((t1 for _, t1 in spans), default=0.0)


# ── scoring ────────────────────────────────────────────────────────────────

def score(pred: Prediction, gt, points: list[float], mask: set[str] | None,
          count_out_of_cast_fp: bool = False):
    """Micro-sum TP/FP/FN over timepoints.

    Each side is a set of actors, an actor being its key-set. Predicted actor P
    matches GT actor G iff their key-sets intersect (any shared id/name). We match
    greedily so each actor is used once, then:
      TP = matched pairs, FP = unmatched predicted, FN = unmatched GT.

    `mask` (gallery∩GT keys) restricts GT so X-Ray cast we can't recognise doesn't
    inflate FN. By default predictions are masked the same way — which DROPS a
    predicted actor who isn't in this film's cast (a cross-film misidentification),
    hiding the pipeline's worst false positives.

    Set count_out_of_cast_fp=True to keep ALL predictions: an actor named who is not
    a present GT cast member counts as an FP, including out-of-cast confusions. This
    is the honest, ship-relevant precision. GT is still masked for fair recall.
    """
    TP = FP = FN = 0
    per_point = []

    for t in points:
        P = [set(a) for a in pred.present_at(t)]
        G = [set(a) for a in gt.present_at(t)]
        if mask is not None:
            G = [a for a in G if a & mask]
            if not count_out_of_cast_fp:
                P = [a for a in P if a & mask]

        tp = _match_count(P, G)
        fp = len(P) - tp
        fn = len(G) - tp
        TP += tp
        FP += fp
        FN += fn
        per_point.append((t, tp, fp, fn))

    prec = TP / (TP + FP) if (TP + FP) else 0.0
    rec = TP / (TP + FN) if (TP + FN) else 0.0
    f1 = 2 * prec * rec / (prec + rec) if (prec + rec) else 0.0
    return {"TP": TP, "FP": FP, "FN": FN, "precision": prec,
            "recall": rec, "f1": f1, "n_points": len(points),
            "per_point": per_point}


def _match_count(P: list[set[str]], G: list[set[str]]) -> int:
    """Greedy 1:1 matching of predicted↔GT actors by key intersection."""
    used = [False] * len(G)
    matched = 0
    for pa in P:
        for j, ga in enumerate(G):
            if not used[j] and pa & ga:
                used[j] = True
                matched += 1
                break
    return matched


# ── main ────────────────────────────────────────────────────────────────────

def main():
    p = argparse.ArgumentParser(description=__doc__,
                                formatter_class=argparse.RawDescriptionHelpFormatter)
    p.add_argument("--pred", required=True, help="pipeline output JSON")
    src = p.add_mutually_exclusive_group(required=True)
    src.add_argument("--xray", help="dir with people.csv/scenes.csv/people_in_scenes.csv")
    src.add_argument("--movienet", help="MovieNet-PS root (needs --split and --title)")
    p.add_argument("--split", default="Train_app10", help="MovieNet annotation split")
    p.add_argument("--title", help="MovieNet title tt-id to filter to")
    p.add_argument("--gallery", help="gallery.json for masking (gallery ∩ GT)")
    p.add_argument("--crosswalk", help="tmdb→imdb JSON (tmdb_imdb_map.py) for exact "
                                       "id join when pred/gallery lack imdb_id")
    p.add_argument("--no-mask", action="store_true", help="disable gallery∩GT masking")
    p.add_argument("--step", type=float, default=1.0, help="regular grid step (s)")
    p.add_argument("--random", type=int, help="sample N uniform-random timepoints")
    p.add_argument("--scene-anchored", action="store_true",
                   help="sample GT scene midpoints (one vote per scene)")
    p.add_argument("--end", type=float, help="timeline end (s); default = max of pred/GT")
    p.add_argument("--seed", type=int, default=0)
    p.add_argument("--json-out", help="write full metrics (incl. per-point) here")
    args = p.parse_args()

    crosswalk = CrosswalkTable.load(args.crosswalk) if args.crosswalk else None
    if crosswalk is not None:
        print(f"[eval] crosswalk: {len(crosswalk)} tmdb→imdb entries", file=sys.stderr)

    pred = Prediction(args.pred, crosswalk=crosswalk)
    print(f"[eval] pred: {len(pred.actors)} actors, timeline≈{pred.max_t:.0f}s "
          f"({pred.movie})", file=sys.stderr)

    if args.xray:
        gt = XRayGroundTruth(args.xray)
    else:
        if not args.title:
            sys.exit("[eval] --movienet requires --title tt-id")
        gt = _load_movienet(args.movienet, args.split, args.title, args.gallery)
    print(f"[eval] GT: {gt.summary()}", file=sys.stderr)

    mask = None
    if not args.no_mask:
        gkeys = load_gallery_keys(args.gallery, crosswalk=crosswalk)
        gt_keys = gt.all_keys()
        if gkeys is None:
            # no gallery given → mask to GT ∩ pred key spaces so absent-from-gallery
            # GT actors don't inflate FN. Fall back to GT keys the pred could name.
            mask = gt_keys & pred.all_keys()
            print("[eval] no --gallery; masking to GT∩pred keys "
                  f"({len(mask)})", file=sys.stderr)
        else:
            mask = gkeys & gt_keys
            print(f"[eval] mask = gallery∩GT ({len(mask)} keys)", file=sys.stderr)

    points = sample_points(args, pred, gt)
    print(f"[eval] sampling {len(points)} timepoints "
          f"({'scene-anchored' if args.scene_anchored else 'random' if args.random else f'grid@{args.step}s'})",
          file=sys.stderr)

    m = score(pred, gt, points, mask)
    print("\n── presence eval ─────────────────────────────")
    print(f"  timepoints : {m['n_points']}")
    print(f"  TP/FP/FN   : {m['TP']} / {m['FP']} / {m['FN']}")
    print(f"  precision  : {m['precision']*100:.1f}%")
    print(f"  recall     : {m['recall']*100:.1f}%")
    print(f"  F1         : {m['f1']*100:.1f}%")

    if args.json_out:
        out = {k: v for k, v in m.items() if k != "per_point"}
        out["per_point"] = [{"t": t, "tp": tp, "fp": fp, "fn": fn}
                            for t, tp, fp, fn in m["per_point"]]
        Path(args.json_out).write_text(json.dumps(out, indent=2))
        print(f"[eval] wrote {args.json_out}", file=sys.stderr)


def _load_movienet(root, split, title, gallery):
    sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
    from movienet_prep import load_movienet_annotations
    anns = load_movienet_annotations(Path(root), split)
    anns = [a for a in anns if a["img_path"].startswith(title)]
    if not anns:
        sys.exit(f"[eval] no MovieNet annotations for title {title} in {split}")
    id_to_name = {}
    if gallery:
        for a in json.load(open(gallery)).get("actors", []):
            if a.get("imdb_id"):
                id_to_name[a["imdb_id"]] = a.get("name", "")
    return MovieNetGroundTruth(anns, id_to_name)


if __name__ == "__main__":
    main()