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

#!/usr/bin/env python3
"""
Self-contained tests for the presence eval. Builds a synthetic X-Ray fixture and
a pipeline-output JSON in a temp dir, then checks scoring, masking, name-fallback
matching, and sampling modes. Run: python scripts/validation/test_sample_eval.py
"""

import csv
import json
import sys
import tempfile
from pathlib import Path

HERE = Path(__file__).resolve().parent
sys.path.insert(0, str(HERE))

from identity import keys_for, norm_name          # noqa: E402
from sample_eval import Prediction, score, sample_points, _match_count  # noqa: E402
from ground_truth import XRayGroundTruth           # noqa: E402
from tmdb_imdb_map import CrosswalkTable           # noqa: E402


def _write_xray(d: Path):
    d.mkdir(parents=True, exist_ok=True)
    with open(d / "scenes.csv", "w", newline="") as f:
        w = csv.writer(f); w.writerow(["scene", "start_ms", "end_ms"])
        w.writerows([("1", 0, 60000), ("2", 300000, 340000)])
    # real Zenodo X-Ray schema: people(name_id,person,character);
    # people_in_scenes(scene,start,end,name_id,timestamp)
    with open(d / "people.csv", "w", newline="") as f:
        w = csv.writer(f); w.writerow(["name_id", "person", "character"])
        w.writerows([("nm0330687", "Lauren Graham", "Lorelai"),
                     ("nm0004754", "Alexis Bledel", "Rory"),
                     ("nm0000001", "Ghost Actor",   "Ghost")])
    with open(d / "people_in_scenes.csv", "w", newline="") as f:
        w = csv.writer(f); w.writerow(["scene", "start", "end", "name_id", "timestamp"])
        w.writerows([("1", 0, 60000, "nm0330687", 5000),
                     ("1", 0, 60000, "nm0000001", 8000),
                     ("2", 300000, 340000, "nm0004754", 305000)])


def _write_pred(path: Path):
    # pipeline output carries tmdb+name but NO nm ids -> name-fallback join to X-Ray
    doc = {"schema_version": 1, "movie": "x", "sample_fps": 1.0, "anneal_sec": 10.0,
           "actors": [
               {"name": "Lauren Graham", "imdb_id": "", "tmdb_id": "16858",
                "jellyfin_id": "a", "scenes": [[21.0, 31.0], [50.0, 59.0]]},
               {"name": "Alexis Bledel", "imdb_id": "", "tmdb_id": "6279",
                "jellyfin_id": "b", "scenes": [[301.0, 311.0]]},
               {"name": "Edward Herrmann", "imdb_id": "", "tmdb_id": "52995",
                "jellyfin_id": "c", "scenes": [[4.0, 56.0]]}]}
    path.write_text(json.dumps(doc))


class T:
    n = 0
    def check(self, cond, msg):
        T.n += 1
        assert cond, f"FAIL: {msg}"
        print(f"  ok: {msg}")


def main():
    t = T()

    # -- identity ---------------------------------------------------------------
    t.check(norm_name("Zöe Saldaña!") == "zoe saldana", "accent/punct normalization")
    t.check(keys_for(imdb_id="nm1", name="Jo Ann") == {"imdb:nm1", "name:jo ann"},
            "keys_for builds namespaced tokens")
    t.check(keys_for(imdb_id="") == set(), "empty ids dropped")

    # -- crosswalk: tmdb-only actor gains an exact imdb: key --------------------
    xw = CrosswalkTable({"16858": "nm0330687", "999": None})
    k = keys_for(imdb_id="", tmdb_id="16858", name="Lauren Graham", crosswalk=xw)
    t.check("imdb:nm0330687" in k, "crosswalk resolves tmdb→imdb key")
    k_null = keys_for(tmdb_id="999", crosswalk=xw)
    t.check(not any(x.startswith("imdb:") for x in k_null), "crosswalk null → no imdb key")
    t.check(len(xw) == 1, "CrosswalkTable len counts non-null entries")

    # -- matching ---------------------------------------------------------------
    t.check(_match_count([{"name:jo"}, {"name:al"}], [{"imdb:x", "name:jo"}]) == 1,
            "one match by shared name key")
    t.check(_match_count([{"name:jo"}], [{"name:jo"}, {"name:jo"}]) == 1,
            "greedy 1:1 uses each GT once")

    with tempfile.TemporaryDirectory() as tmp:
        tmp = Path(tmp)
        _write_xray(tmp / "xray")
        _write_pred(tmp / "pred.json")

        pred = Prediction(tmp / "pred.json")
        gt = XRayGroundTruth(tmp / "xray")

        # -- presence lookups ---------------------------------------------------
        t.check(len(pred.present_at(25)) == 2, "Graham+Herrmann present at 25s")
        t.check(len(pred.present_at(305)) == 1, "only Bledel present at 305s")
        t.check(len(gt.present_at(30)) == 2, "X-Ray scene1 has Graham+Ghost at 30s")
        t.check(len(gt.present_at(320)) == 1, "X-Ray scene2 has Bledel at 320s")

        # -- masking behaviour --------------------------------------------------
        pts = [30.0]
        # unmasked: at 30s pred={Graham,Herrmann}, gt={Graham,Ghost}
        #   match Graham -> TP1; Herrmann unmatched -> FP1; Ghost unmatched -> FN1
        m = score(pred, gt, pts, mask=None)
        t.check((m["TP"], m["FP"], m["FN"]) == (1, 1, 1), "unmasked 30s = 1/1/1")

        # masked to GT∩pred keys: Ghost & Herrmann are absent from the other side's
        # key space, so both drop -> only Graham remains on both -> 1/0/0
        gt_keys = gt.all_keys(); pred_keys = pred.all_keys()
        mask = gt_keys & pred_keys
        m2 = score(pred, gt, pts, mask=mask)
        t.check((m2["TP"], m2["FP"], m2["FN"]) == (1, 0, 0),
                "masked 30s drops off-gallery actors = 1/0/0")

        # -- crosswalk end-to-end: exact imdb join, names garbled ---------------
        # Rebuild pred with names that WON'T match X-Ray, but a crosswalk that maps
        # their tmdb ids to the correct nm ids. Match must survive via imdb key.
        garbled = {"schema_version": 1, "movie": "x", "actors": [
            {"name": "WRONG NAME A", "imdb_id": "", "tmdb_id": "16858",
             "jellyfin_id": "a", "scenes": [[21.0, 31.0]]},            # →nm0330687 Graham
            {"name": "WRONG NAME B", "imdb_id": "", "tmdb_id": "6279",
             "jellyfin_id": "b", "scenes": [[301.0, 311.0]]}]}          # →nm0004754 Bledel
        (tmp / "garbled.json").write_text(json.dumps(garbled))
        xw = CrosswalkTable({"16858": "nm0330687", "6279": "nm0004754"})
        pred_g = Prediction(tmp / "garbled.json", crosswalk=xw)
        # at 30s (scene1) Graham should match by imdb despite wrong name
        m_g = score(pred_g, gt, [30.0], mask=None)
        t.check(m_g["TP"] == 1, "crosswalk yields exact imdb match despite wrong names")
        # without crosswalk, wrong names → no match at all
        pred_bad = Prediction(tmp / "garbled.json")
        m_bad = score(pred_bad, gt, [30.0], mask=None)
        t.check(m_bad["TP"] == 0, "no crosswalk + wrong names → no match")

        # -- sampling modes -----------------------------------------------------
        class A:  # arg stub
            scene_anchored = True; random = None; step = 1.0; end = None; seed = 0
        sp = sample_points(A, pred, gt)
        t.check(sp == [30.0, 320.0], "scene-anchored samples scene midpoints")

        A.scene_anchored = False; A.end = 10.0; A.step = 2.0
        grid = sample_points(A, pred, gt)
        t.check(grid == [0.0, 2.0, 4.0, 6.0, 8.0, 10.0], "regular grid step")

    print(f"\nALL {T.n} CHECKS PASSED")


if __name__ == "__main__":
    main()