scene-actor-extraction/scripts/optimizer/dump_error_frames.py

#!/usr/bin/env python3
"""
dump_error_frames.py — extract example video frames for visual inspection of a
replayed prediction vs X-Ray ground truth: best-agreement seconds, FPI (false
identification) seconds, and FN (missed cast) seconds.

Reuses second_score.py's per-second timeline/prediction loading, but keeps the
per-second classification (score_seconds only returns aggregates) and picks
representative timestamps in each bucket, then pulls single frames from the
source video via ffmpeg -ss (nearest keyframe-independent seek + decode).

If --raw (the JSONL from `replay.py --raw-out`) is given, also draws each visible
actor's bounding box + name/similarity on the extracted frame — green for
identified, orange for unknown — matching debug_renderer_node.hpp's colour
convention. Without --raw, frames are saved unannotated.

Usage:
  python scripts/optimizer/dump_error_frames.py \
      --pred pred.json --raw raw.jsonl \
      --xray experiments/xray/.../900_The_Many_Saints_Of_Newark \
      --movie "/mnt/movies/The Many Saints Of Newark (2021)/....mp4" \
      --gallery experiments/galleries/gallery_LVFace-B_Glint360K.h5 \
      --out-dir experiments/dump_review/many_saints --n-per-bucket 6
"""
from __future__ import annotations

import argparse
import json
import subprocess
import sys
from pathlib import Path

import cv2

REPO = Path(__file__).resolve().parent.parent.parent
sys.path.insert(0, str(REPO / "scripts" / "optimizer"))
sys.path.insert(0, str(REPO / "scripts" / "validation"))

from second_score import load_second_timeline, load_pred_intervals, _match  # noqa: E402
from sample_eval import load_gallery_keys  # noqa: E402
from identity import keys_for  # noqa: E402


def per_second_detail(pred_json: dict, xray_dir: str, gallery_keys: set | None):
    """Like second_score.score_seconds, but yields one record per sampled second
    instead of collapsing to aggregates."""
    timeline, film_cast, duration = load_second_timeline(xray_dir)
    pred = load_pred_intervals(pred_json)

    name_by_keys = {}
    for a in pred_json.get("actors", []):
        k = frozenset(keys_for(imdb_id=a.get("imdb_id"), tmdb_id=a.get("tmdb_id"),
                                jellyfin_id=a.get("jellyfin_id"), name=a.get("name")))
        name_by_keys[k] = a.get("name", "?")

    records = []
    for t in sorted(timeline):
        G = [set(a) for a in timeline[t]]
        P_all = [(k, set(k)) for k, wins in pred if any(w0 <= t <= w1 for w0, w1 in wins)]
        if gallery_keys is not None:
            G = [g for g in G if g & gallery_keys]

        P = [p for _, p in P_all]
        tp, matched = _match(P, G)
        fp_names, fn_names = [], []
        for key, pa in P_all:
            if not any(pa & ga for ga in G):
                fp_names.append(name_by_keys.get(key, "?"))
        for j, ga in enumerate(G):
            if not matched[j]:
                fn_names.append("|".join(sorted(x for x in ga if not x.startswith("imdb:") and not x.startswith("tmdb:"))) or "?")

        union = tp + len(fp_names) + len(fn_names)
        jaccard = (tp / union) if union else 1.0
        records.append({"t": t, "tp": tp, "fp": fp_names, "fn": fn_names, "jaccard": jaccard})
    return records


def pick_timestamps(records, n_per_bucket):
    best = sorted(records, key=lambda r: (-r["jaccard"], -r["tp"]))
    best = [r for r in best if r["tp"] > 0][:n_per_bucket]

    fpi = [r for r in records if r["fp"]]
    fpi = sorted(fpi, key=lambda r: -len(r["fp"]))[:n_per_bucket]

    fn = [r for r in records if r["fn"]]
    fn = sorted(fn, key=lambda r: -len(r["fn"]))[:n_per_bucket]

    return {"best": best, "fpi": fpi, "fn": fn}


def pick_by_interval(records, interval_sec):
    """One best (highest jaccard) and one worst (lowest jaccard) second per
    interval_sec-second window across the whole film, e.g. --interval-sec 600 for
    a per-10-minute best/worst sweep. Windows with no sampled seconds are skipped
    (X-Ray timelines only cover scenes, so gaps between/after scenes are common)."""
    windows: dict[int, list] = {}
    for r in records:
        windows.setdefault(r["t"] // interval_sec, []).append(r)

    buckets: dict[str, list] = {}
    for w in sorted(windows):
        wr = windows[w]
        best = max(wr, key=lambda r: (r["jaccard"], r["tp"]))
        worst = min(wr, key=lambda r: (r["jaccard"], -max(len(r["fp"]), len(r["fn"]))))
        buckets[f"w{w:03d}_best"] = [best]
        buckets[f"w{w:03d}_worst"] = [worst]
    return buckets


def load_raw_annotations(raw_path: str):
    """second (int, floor) -> list of visible_actors dicts (last frame wins if
    several fall in the same second, which is the common case at 1fps sampling)."""
    by_second = {}
    with open(raw_path) as f:
        for line in f:
            sa = json.loads(line)
            if sa.get("eof"):
                continue
            by_second[int(sa["timestamp_sec"])] = sa.get("visible_actors", [])
    return by_second


def draw_annotations(frame_path: Path, actors: list):
    img = cv2.imread(str(frame_path))
    if img is None:
        return
    for a in actors:
        known = a.get("actor_idx", -1) >= 0
        colour = (60, 200, 0) if known else (220, 100, 0)  # BGR: green / orange
        x, y, w, h = a["bbox"]
        x, y, w, h = int(x), int(y), int(w), int(h)
        cv2.rectangle(img, (x, y), (x + w, y + h), colour, 2)

        label = f"{a['name']} {a['similarity']*100:.0f}%" if known else f"unknown {a['similarity']*100:.0f}%"
        (tw, th), baseline = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)
        strip_y0 = max(0, y - th - 4)
        cv2.rectangle(img, (x, strip_y0), (x + tw + 4, y), colour, cv2.FILLED)
        cv2.putText(img, label, (x + 2, y - 2), cv2.FONT_HERSHEY_SIMPLEX, 0.5,
                    (255, 255, 255), 1, cv2.LINE_AA)
    cv2.imwrite(str(frame_path), img)


def extract_frame(movie: str, t: float, out_path: Path):
    out_path.parent.mkdir(parents=True, exist_ok=True)
    subprocess.run(
        ["ffmpeg", "-y", "-ss", str(t), "-i", movie, "-frames:v", "1",
         "-q:v", "2", str(out_path)],
        check=True, capture_output=True)


def main():
    p = argparse.ArgumentParser(description=__doc__,
                                formatter_class=argparse.RawDescriptionHelpFormatter)
    p.add_argument("--pred", required=True)
    p.add_argument("--raw", help="raw per-frame annotations JSONL (replay.py --raw-out); "
                                  "draws bboxes + names on extracted frames if given")
    p.add_argument("--xray", required=True)
    p.add_argument("--movie", required=True)
    p.add_argument("--gallery")
    p.add_argument("--out-dir", required=True)
    p.add_argument("--n-per-bucket", type=int, default=6)
    p.add_argument("--interval-sec", type=int,
                   help="instead of global best/fpi/fn buckets, pick one best + one "
                        "worst (by jaccard) second per interval-sec window across "
                        "the whole film, e.g. 600 for per-10-minute best/worst")
    args = p.parse_args()

    pred_json = json.loads(Path(args.pred).read_text())
    gk = load_gallery_keys(args.gallery) if args.gallery else None
    records = per_second_detail(pred_json, args.xray, gk)
    buckets = (pick_by_interval(records, args.interval_sec) if args.interval_sec
              else pick_timestamps(records, args.n_per_bucket))
    raw_by_second = load_raw_annotations(args.raw) if args.raw else None

    out_dir = Path(args.out_dir)
    manifest = []
    for bucket, recs in buckets.items():
        for r in recs:
            fname = f"{bucket}_t{r['t']:05d}.jpg"
            out_path = out_dir / bucket / fname
            try:
                extract_frame(args.movie, r["t"], out_path)
                ok = True
                if raw_by_second is not None:
                    draw_annotations(out_path, raw_by_second.get(r["t"], []))
            except subprocess.CalledProcessError as e:
                ok = False
                print(f"[dump_error_frames] ffmpeg failed at t={r['t']}: {e}", file=sys.stderr)
            manifest.append({"bucket": bucket, "t": r["t"], "tp": r["tp"],
                             "fp": r["fp"], "fn": r["fn"], "jaccard": round(r["jaccard"], 3),
                             "file": str(out_path.relative_to(out_dir)) if ok else None})
            print(f"[{bucket}] t={r['t']}s tp={r['tp']} fp={r['fp']} fn={r['fn']}", file=sys.stderr)

    (out_dir / "manifest.json").write_text(json.dumps(manifest, indent=2, ensure_ascii=False))
    print(f"[dump_error_frames] wrote {len(manifest)} frames + manifest.json to {out_dir}",
          file=sys.stderr)


if __name__ == "__main__":
    main()