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

#!/usr/bin/env python3
"""
replay.py — replay a dumped embedding HDF5 through the real KPN downstream nodes.

TRACES: VR-002, VR-011 | PR-002

Reads an embedding dump (scripts/optimizer/SCHEMA.md), feeds each frame as an
EmbeddedSceneFrame into a Python-assembled KPN network wiring the *real* C++
face_tracker → identity_matcher → frame_annotation → result_sink, and reads back
the truth file that sink wrote. No decode, no GPU embedding — only the cheap
downstream tail runs, so a sweep can vary Config knobs freely.

The sink is part of the network, not a Python reimplementation of it. That is
VR-011: presence comes from TrackRegistry claims, so a replayed window and a
scene_analyze window are produced by the same code rather than by two functions
that agreed once. See [[kpn-python-replay-optimizer]].

CLI:
  python scripts/optimizer/replay.py --dump film.h5 --gallery gallery.json \
      --out replayed.json [--prob-threshold 0.99] [--track-extinction-sec 5] ...
"""
from __future__ import annotations

import argparse
import json
import sys
import time
from pathlib import Path

import h5py
import numpy as np

REPO = Path(__file__).resolve().parent.parent.parent
sys.path.insert(0, str(REPO / "scripts"))
from sae_stamp import verify_gallery_stamp  # noqa: E402


def dump_embedder_stamp(dump_path: str) -> dict:
    """The GR-004 embedder stamp recorded in an embedding dump.

    A replay has no live embedder — the dump IS the embedder as far as the gallery
    is concerned, so the dump's stamp is what the gallery must be checked against.
    Dumps written before GR-004 have no attributes and yield an empty stamp, which
    the check reports as unverifiable rather than silently accepting."""
    with h5py.File(dump_path, "r") as f:
        name = f.attrs.get("embedder_model", "")
        sha = f.attrs.get("embedder_sha256", "")
    dec = lambda v: v.decode() if isinstance(v, bytes) else ("" if v is None else str(v))
    return {"model_name": dec(name), "model_sha256": dec(sha), "embed_dim": 512}


def load_frames(dump_path: str, min_conf: float = 0.0):
    """Yield EmbeddedSceneFrame dicts from the HDF5 dump, then a trailing EOF.

    `min_conf` drops detections below that detector confidence before they reach the
    matcher — an UPWARD-only detector_conf sweep on already-dumped faces (the dump was
    made at detector_conf=0.5, so 0.5 is the floor). Lets us test whether near-threshold
    detections are real faces (raising min_conf hurts recall) or phantoms (it helps
    precision at no recall cost)."""
    with h5py.File(dump_path, "r") as f:
        ts   = f["frames/timestamp_sec"][:]
        fidx = f["frames/frame_idx"][:]
        cut  = f["frames/is_cut"][:]
        off  = f["frames/face_offset"][:]
        cnt  = f["frames/face_count"][:]
        emb  = f["faces/embedding"][:]
        bbox = f["faces/bbox"][:]
        lmk  = f["faces/landmarks"][:]
        conf = f["faces/confidence"][:]
        # TRACES: AR-028 | SR-002
        # The quality vector, present from schema v2. A v1 dump predates AR-028
        # and simply has no such dataset — read as absent, never as a default,
        # so a face from an old dump stays at the C++ -1 "unscored" sentinel
        # rather than acquiring a fabricated sharpness of 0 (which is a real
        # value on this axis, meaning a featureless crop).
        qual = {k: f[f"faces/{k}"][:] for k in ("sharpness", "alignment_residual")
                if f"faces/{k}" in f}
        movie = f.attrs.get("movie", "")
        fps   = float(f.attrs.get("sample_fps", 1.0))

    frames = []
    for i in range(len(ts)):
        s, n = int(off[i]), int(cnt[i])
        keep = slice(s, s + n)
        c = np.ascontiguousarray(conf[keep], dtype=np.float32)
        if min_conf > 0.0 and n:
            m = c >= min_conf
            sel = np.where(m)[0]
            frames.append({
                "timestamp_sec": float(ts[i]), "frame_idx": int(fidx[i]),
                "is_cut": bool(cut[i]), "eof": False,
                "bbox":       np.ascontiguousarray(bbox[keep][sel], dtype=np.float32),
                "landmarks":  np.ascontiguousarray(lmk[keep][sel],  dtype=np.float32),
                "confidence": np.ascontiguousarray(c[sel],          dtype=np.float32),
                "embeddings": np.ascontiguousarray(emb[keep][sel],  dtype=np.float32),
                **{k: np.ascontiguousarray(v[keep][sel], dtype=np.float32)
                   for k, v in qual.items()},
            })
        else:
            frames.append({
                "timestamp_sec": float(ts[i]), "frame_idx": int(fidx[i]),
                "is_cut": bool(cut[i]), "eof": False,
                "bbox":       np.ascontiguousarray(bbox[keep], dtype=np.float32),
                "landmarks":  np.ascontiguousarray(lmk[keep],  dtype=np.float32),
                "confidence": c,
                "embeddings": np.ascontiguousarray(emb[keep],  dtype=np.float32),
                **{k: np.ascontiguousarray(v[keep], dtype=np.float32)
                   for k, v in qual.items()},
            })
    last_ts = float(ts[-1]) if len(ts) else 0.0
    frames.append({"timestamp_sec": last_ts, "eof": True})
    return frames, str(movie), fps


def replay(dump_path: str, gallery: str, cfg: dict, build_dir: str,
           out_path: str, stop: bool = True, raw_out: str | None = None,
           eof_timeout: float = 300.0) -> dict:
    """Run the dump through the real KPN chain and return the truth file it wrote.

    TRACES: VR-011, VR-002 | PR-002

    `out_path` is where the C++ sink writes. That is the change VR-011 makes:
    the presence windows in that file are built by ResultSinkFunc from
    TrackRegistry claims -- the extent of a track an actor owned (AR-012),
    ending at the last sighting (AR-013) -- and are byte-for-byte the same
    construction scene_analyze ships. This function used to build them itself,
    in Python, by annealing gaps between per-frame detections, which is what the
    pipeline did BEFORE AR-012. A sweep tuned against that was tuning a contract
    the shipped code had stopped honouring.

    cfg may include "detector_conf" to prune dumped detections below that
    confidence (upward-only from the 0.5 dump floor) before matching.

    raw_out: if set, also write per-frame annotations as JSON lines for the
    montage renderers. Derived from the truth file's own `frames` array rather
    than tapped separately out of the network -- see write_raw_frames.

    eof_timeout: how long to wait for the sink to write. A replay that never
    reaches EOF is a wedged pipeline, and returning an empty result would look
    like a film with no cast rather than like a failure."""
    sys.path.insert(0, build_dir)
    import sae_kpn

    # TRACES: GR-004 | SR-001
    # checked here, before any network is built, so a
    # cross-model replay dies with one readable error instead of producing a
    # plausible-looking score. add_identity_matcher re-checks it C++-side below;
    # that is the backstop for any other caller of the binding.
    stamp = dump_embedder_stamp(dump_path)
    verify_gallery_stamp(gallery, stamp=stamp,
                         embedder_desc=f"embedding dump {Path(dump_path).name}",
                         require_stamp=bool(cfg.get("require_gallery_stamp", False)))

    frames, movie, fps = load_frames(dump_path, min_conf=float(cfg.get("detector_conf", 0.0)))

    net = sae_kpn.Network()
    sae_kpn._register_types(net)

    idx = [0]
    eof = {"timestamp_sec": frames[-1]["timestamp_sec"], "eof": True}

    def source():
        # A no-input source node's run_loop calls this in a tight loop. Once frames
        # are exhausted we must NOT hot-spin returning EOF — that pegs a core and
        # floods the downstream channel with EOFs (livelock that wedged DE). Sleep
        # briefly after the single real EOF so net.stop() can tear the thread down.
        i = idx[0]
        idx[0] += 1
        if i < len(frames):
            return frames[i]
        time.sleep(0.05)
        return eof

    # Channel capacity must exceed the frame count so the fast source can't overflow
    # a downstream FIFO before the serial reader drains it — PyNode DROPS on overflow,
    # which would silently truncate the replay. Size to the whole film + slack.
    # Every channel gets capacity ≥ the whole film so NOTHING can ever overflow-drop:
    # the source can push all frames before any downstream node has drained, and a
    # dropped frame silently corrupts the score. Memory is cheap (a few k pointers);
    # correctness is not. Generous slack on top.
    cap = len(frames) * 2 + 64
    sae_kpn.add_node_python(net, "replay", source, [], ["EmbeddedSceneFrame"], cap)

    # TRACES: VR-011, VR-002 | DP-001 | PR-002
    # One call builds tracker -> matcher -> annotation -> sink in the only order
    # that works (the matcher fits the calibration the tracker needs, and the
    # sink needs the registry's claims). This used to be three factory calls
    # assembled here, which is how the seam broke: the ordering constraint could
    # not be expressed, so the tracker was built from a Config alone long after
    # it had started requiring a registry and a calibration.
    cfg = dict(cfg)
    cfg["output_path"] = out_path
    cfg["movie_path"]  = movie
    cfg["sample_fps"]  = fps
    # Verbosity 1 (standard) adds the per-frame array; only pay for it when the
    # caller wants raw frames, since it retains every annotation in memory.
    cfg["verbosity"] = 1 if raw_out else 0
    sae_kpn.add_pipeline(net, gallery, cfg, cap,
                         stamp["model_name"], stamp["model_sha256"])

    net.connect("replay", 0, "tracker", 0)
    net.connect("tracker", 0, "matcher", 0)
    net.connect("matcher", 0, "annotation", 0)
    net.connect("annotation", 0, "sink", 0)
    net.build()
    net.start()

    # The sink writes on the EOF annotation. Wait for it rather than reading
    # anything back through the seam: presence is the registry's answer, and the
    # registry lives entirely on the C++ side.
    #
    # This replaces a read loop that pulled one SceneAnnotation per input frame
    # and rebuilt windows in Python. That loop needed a heuristic -- "keep
    # reading past eof until we've collected all n_frames annotations, or hit a
    # run of 8 consecutive eofs" -- to work around a tail it was losing. None of
    # that exists now: nothing is read per frame, so nothing can be lost per
    # frame.
    deadline = time.time() + eof_timeout
    while not sae_kpn.pipeline_done(net):
        if time.time() > deadline:
            sae_kpn.release_pipeline(net)
            raise TimeoutError(
                f"replay did not finish within {eof_timeout}s "
                f"({len(frames) - 1} frames); the sink never saw EOF")
        time.sleep(0.02)

    if stop:
        net.stop()
    sae_kpn.release_pipeline(net)

    with open(out_path) as f:
        result = json.load(f)

    if raw_out:
        write_raw_frames(result, raw_out)
    return result


def write_raw_frames(truth: dict, raw_out: str) -> None:
    """Per-frame annotations as JSONL, for the montage/error-frame renderers.

    TRACES: VR-011 | PR-002

    Derived from the truth file's own `frames` array (verbosity 1) rather than
    from a second stream tapped out of the network. One producer, one set of
    numbers: a bbox drawn on a montage is now provably the bbox the sink
    recorded, which it was not when Python read annotations separately.

    The shape is the legacy one -- {timestamp_sec, visible_actors:[...]} with
    actor_idx/bbox/name/similarity -- because dump_scene_montage.py and
    dump_error_frames.py read exactly those fields, and rewriting them is not
    what this requirement is about.
    """
    with open(raw_out, "w") as f:
        for fr in truth.get("frames", []):
            visible = []
            for a in fr.get("identified", []):
                visible.append({
                    "actor_idx":  0,          # >= 0 means "known"; the renderers
                                              # test the sign, never the value
                    "name":       a.get("name", ""),
                    "imdb_id":    a.get("imdb_id", ""),
                    "tmdb_id":    a.get("tmdb_id", ""),
                    "jellyfin_id": a.get("jellyfin_id", ""),
                    "similarity": a.get("similarity", 0.0),
                    "track_id":   a.get("track_id", -1),
                    "bbox":       a.get("bbox", [0, 0, 0, 0]),
                })
            for u in fr.get("unknowns", []):
                visible.append({
                    "actor_idx":  -1,
                    "name":       "",
                    "similarity": u.get("confidence", 0.0),
                    "track_id":   u.get("track_id", -1),
                    "bbox":       u.get("bbox", [0, 0, 0, 0]),
                })
            f.write(json.dumps({"timestamp_sec": fr.get("t", 0.0),
                                "visible_actors": visible}) + "\n")


CFG_KEYS = ["detector_conf", "prob_threshold", "match_prior",
            "track_alpha", "track_min_iou", "track_assoc_min_prob",
            "track_extinction_sec",
            # AR-025 ownership and evidence accumulation. Newly reachable:
            # these were in-class defaults no sweep could vary, which is why
            # VR-007 never covered them despite rho_max deferring to it.
            "ownership_logodds", "evidence_rho_max", "evidence_admit_below",
            "evidence_max_views"]

# TRACES: VR-011 | PR-002
# REPLAY_LOCAL_KEYS is gone with build_minimal. It held anneal_sec, the last
# parameter this harness applied itself -- and the only reason it needed a
# separate list was that the harness was still doing windowing the pipeline had
# stopped doing. Every key is a Config key now, because every decision is the
# pipeline's.


def main():
    p = argparse.ArgumentParser(description=__doc__,
                                formatter_class=argparse.RawDescriptionHelpFormatter)
    p.add_argument("--dump", required=True, help="embedding HDF5 dump")
    p.add_argument("--gallery", required=True)
    p.add_argument("--out", required=True, help="output presence JSON")
    p.add_argument("--raw-out", help="also write raw per-frame annotations (JSONL, with bboxes) here")
    p.add_argument("--build-dir", default=str(REPO / "build"))
    for k in CFG_KEYS:
        p.add_argument(f"--{k.replace('_','-')}", type=float, default=None)
    # per-film gallery expansion: promotes pose-varied views of confidently-identified
    # actors into an in-memory annex, recovering ~+4 recall at no precision cost.
    p.add_argument("--expand-gallery", action="store_true")
    # TRACES: GR-004 | SR-001
    # promote an unprovable gallery/dump binding from a
    # loud warning to a hard error. Measurement sweeps should set this (or
    # SAE_REQUIRE_GALLERY_STAMP=1) so no number comes from an unbound pair.
    p.add_argument("--require-gallery-stamp", action="store_true")
    args = p.parse_args()

    cfg = {k: getattr(args, k) for k in CFG_KEYS if getattr(args, k) is not None}
    if args.expand_gallery:
        cfg["expand_gallery"] = True
    if args.require_gallery_stamp:
        cfg["require_gallery_stamp"] = True
    # stop=True: PyNode::stop() sets stop_flag_ before joining, so the source
    # thread's run_loop actually exits. stop=False skips that, leaving stop_flag_
    # false forever — the PyNode destructor's jthread.join() then blocks forever
    # (verified via gdb: stuck in the source node's run_loop, not the GEMM path).
    result = replay(args.dump, args.gallery, cfg, args.build_dir,
                    out_path=args.out, stop=True, raw_out=args.raw_out)
    # NOT rewritten here: the sink already wrote args.out, and that file is the
    # artifact. Dumping `result` back over it would make this script the last
    # writer of a file it did not produce -- and any formatting difference would
    # be a diff between the replayed truth file and a scene_analyze one that is
    # this script's doing rather than the pipeline's.
    print(f"[replay] {len(result['actors'])} actors → {args.out}", file=sys.stderr)


if __name__ == "__main__":
    main()