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

#!/usr/bin/env python3
"""
ground_truth.py — pluggable ground-truth loaders for per-scene presence eval.

A ground truth is a timeline of "who is present when", exposed as:

    GroundTruth.present_at(t: float) -> set[str]      # match-keys present at time t
    GroundTruth.scene_windows()     -> list[(t0, t1)] # scene spans (for --scene-anchored)
    GroundTruth.all_keys()          -> set[str]        # every actor the GT knows (for masking)

"Match-keys" are provider-agnostic identity tokens (see identity.py): an actor is
represented by *all* the keys we can derive (nm-id, tmdb-id, jellyfin-id, normalized
name), so predicted and GT sets intersect if they agree on *any* shared id space.
This matters because the pipeline output may carry only tmdb/jellyfin ids while
X-Ray/MovieNet key on IMDb nm-ids — see [[per-scene-presence-eval-design]].

Two sources implemented:
  * XRayGroundTruth  — Zenodo scene-level Amazon X-Ray CSVs (cast-in-scene).
  * MovieNetGroundTruth — MovieNet-PS per-shot face annotations (on-screen faces).
"""

from __future__ import annotations

import csv
import sys
from bisect import bisect_right
from pathlib import Path

from identity import keys_for


class GroundTruth:
    """Base: a set of actors, each with presence intervals [(t0,t1), ...] in seconds."""

    def __init__(self) -> None:
        # actor_id (any stable local id) -> {"keys": set[str], "intervals": [(t0,t1)]}
        self._actors: dict[str, dict] = {}
        self._scene_spans: list[tuple[float, float]] = []

    # -- construction helpers ------------------------------------------------
    def _add_interval(self, actor_id: str, keys: set[str], t0: float, t1: float) -> None:
        a = self._actors.setdefault(actor_id, {"keys": set(), "intervals": []})
        a["keys"] |= keys
        a["intervals"].append((float(t0), float(t1)))

    def _finalize(self) -> None:
        """Sort intervals and precompute a flat sorted start-array per actor."""
        for a in self._actors.values():
            a["intervals"].sort()
            a["_starts"] = [iv[0] for iv in a["intervals"]]
        self._scene_spans.sort()

    # -- query API -----------------------------------------------------------
    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.values():
            ivs = a["intervals"]
            i = bisect_right(a["_starts"], t)  # first interval starting after t
            # walk back over intervals that started at/before t
            j = i - 1
            while j >= 0:
                t0, t1 = ivs[j]
                if t1 >= t:
                    out.add(frozenset(a["keys"]))
                    break
                # intervals sorted by start; an earlier one could still cover t,
                # but since we only need membership, keep scanning a bounded window.
                j -= 1
                if i - j > 8:  # bound: overlapping intervals per actor are rare
                    break
        return out

    def scene_windows(self) -> list[tuple[float, float]]:
        return self._scene_spans

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

    def summary(self) -> str:
        n_iv = sum(len(a["intervals"]) for a in self._actors.values())
        return (f"{len(self._actors)} actors, {n_iv} intervals, "
                f"{len(self._scene_spans)} scenes")


# ── Amazon X-Ray (Zenodo) ─────────────────────────────────────────────────────

class XRayGroundTruth(GroundTruth):
    """
    Load one movie's X-Ray CSVs (Zenodo DOI 10.5281/zenodo.17659734).
    Real schema (columns are milliseconds):
      people.csv            name_id (nm...), person, character
      scenes.csv            scene, start, end                  (ms)
      people_in_scenes.csv  scene, start, end, name_id, timestamp   (ms)

    Presence = whole scene span for every character listed in that scene.
    Semantics: cast-in-scene (incl. off-camera) — a recall ceiling, not accuracy.
    Columns are resolved case-insensitively so minor variants still load.
    """

    def __init__(self, movie_dir: str | Path) -> None:
        super().__init__()
        d = Path(movie_dir)
        people = _read_csv(d / "people.csv")
        scenes = _read_csv(d / "scenes.csv")
        pis = _read_csv(d / "people_in_scenes.csv")

        # nm-id -> actor name (for building match keys)
        nm_col_p = _find_col(people, "name_id", "nm", "imdb")
        name_col = _find_col(people, "person", "actor")  # actor name lives in "person"
        id_to_name: dict[str, str] = {}
        for row in people:
            nm = (row.get(nm_col_p) or "").strip()
            if nm:
                id_to_name[nm] = (row.get(name_col) or "").strip()

        # scene number -> (t0_sec, t1_sec)
        scene_col_s = _find_col(scenes, "scene")
        start_col = _find_col(scenes, "start")
        end_col = _find_col(scenes, "end")
        span: dict[str, tuple[float, float]] = {}
        for row in scenes:
            sn = (row.get(scene_col_s) or "").strip()
            t0 = _ms_to_sec(row.get(start_col))
            t1 = _ms_to_sec(row.get(end_col))
            if sn and t0 is not None and t1 is not None:
                span[sn] = (t0, t1)
                self._scene_spans.append((t0, t1))

        # scene number -> [nm ids present]
        scene_col_pis = _find_col(pis, "scene")
        nm_col_pis = _find_col(pis, "name_id", "nm", "imdb")
        for row in pis:
            sn = (row.get(scene_col_pis) or "").strip()
            nm = (row.get(nm_col_pis) or "").strip()
            if sn not in span or not nm:
                continue
            t0, t1 = span[sn]
            self._add_interval(nm, keys_for(imdb_id=nm, name=id_to_name.get(nm)), t0, t1)

        self._finalize()


# ── MovieNet-PS ────────────────────────────────────────────────────────────────

class MovieNetGroundTruth(GroundTruth):
    """
    Build presence from MovieNet-PS per-shot face annotations for a single title.

    Input: the flat annotation list produced by movienet_prep.load_movienet_annotations
    filtered to one movie (tt-id), plus a shot->time map. Because MovieNet frames are
    named tt.../shot_XXXX_img_Y.jpg with no absolute timestamp, presence is expressed
    in *shot index* units unless a fps/shot-duration map is supplied. For the sampler
    we therefore sample at shot granularity (one timepoint per annotated shot).

    Semantics: on-screen face presence per shot — like-for-like fair benchmark.
    """

    def __init__(self, annotations: list[dict], id_to_name: dict[str, str] | None = None,
                 shot_seconds: float = 1.0) -> None:
        super().__init__()
        id_to_name = id_to_name or {}
        # group by shot index; each annotated shot becomes a unit interval on a
        # synthetic timeline (shot_index * shot_seconds).
        shots: dict[int, set[str]] = {}
        for ann in annotations:
            shot = _shot_index(ann["img_path"])
            if shot is None:
                continue
            shots.setdefault(shot, set()).add(ann["imdb_id"])

        for shot, nm_ids in shots.items():
            t0 = shot * shot_seconds
            t1 = t0 + shot_seconds
            self._scene_spans.append((t0, t1))
            for nm in nm_ids:
                self._add_interval(nm, keys_for(imdb_id=nm, name=id_to_name.get(nm)), t0, t1)

        self._finalize()


# ── small parsing helpers ──────────────────────────────────────────────────────

def _read_csv(path: Path) -> list[dict]:
    if not path.exists():
        raise FileNotFoundError(f"expected X-Ray CSV not found: {path}")
    with open(path, newline="", encoding="utf-8") as f:
        return list(csv.DictReader(f))


def _find_col(rows: list[dict], *keywords: str) -> str:
    """Return the first column whose lowercased name contains all keywords of any
    single keyword group. We try each keyword in order and accept the first hit."""
    if not rows:
        raise ValueError("empty CSV — cannot resolve columns")
    cols = list(rows[0].keys())
    low = {c: c.lower() for c in cols}
    for kw in keywords:
        for c in cols:
            if kw in low[c]:
                return c
    raise KeyError(f"no column matching {keywords} in {cols}")


def _ms_to_sec(v) -> float | None:
    if v is None or str(v).strip() == "":
        return None
    try:
        return float(v) / 1000.0
    except ValueError:
        return None


def _shot_index(img_path: str) -> int | None:
    # tt0032138/shot_0003_img_1.jpg -> 3
    import re
    m = re.search(r"shot_(\d+)", img_path)
    return int(m.group(1)) if m else None