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

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

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 → scene_tracker, and returns the same presence-window
JSON that scene_analyze's result_sink produces (minimal schema). No decode, no GPU
embedding — only the cheap downstream tail runs, so a sweep can vary Config knobs
freely. See [[kpn-python-replay-optimizer]].

CLI:
  python scripts/optimizer/replay.py --dump film.h5 --gallery gallery.json \
      --out replayed.json [--prob-threshold 0.99] [--anneal 10] ...
"""
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


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"][:]
        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),
            })
        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),
            })
    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, stop: bool = True,
          raw_out: str | None = None) -> dict:
    """Run the dump through the real KPN chain; return minimal-schema presence JSON.

    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 the raw per-frame annotations (timestamp, actor_idx,
    name, bbox, similarity — one entry per input frame, before merging into windows)
    as JSON lines to this path. Needed to draw bounding boxes on extracted frames;
    the merged window schema returned by this function has no per-frame bbox."""
    sys.path.insert(0, build_dir)
    import sae_kpn

    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)
    sae_kpn.add_face_tracker(net, "tracker", cfg, cap)
    sae_kpn.add_identity_matcher(net, "matcher", gallery, cfg, cap)
    sae_kpn.add_scene_tracker(net, "scene", cfg, cap)
    net.connect("replay", 0, "tracker", 0)
    net.connect("tracker", 0, "matcher", 0)
    net.connect("matcher", 0, "scene", 0)
    net.build()
    net.start()

    # Read exactly one annotation per input frame. The source emits EOF as an ordinary
    # value AFTER the last frame, but the concurrent pipeline lets that EOF OVERTAKE
    # the last few real frames still flowing tracker→matcher→scene. Breaking on the
    # first eof therefore dropped a random tail (~0.5–1%, race-dependent). Instead we
    # keep reading past eof until we've collected all n_frames annotations (or hit a
    # run of consecutive eofs meaning the pipeline is genuinely drained).
    n_expected = len(frames) - 1   # excludes the trailing eof frame
    annotations = []
    eof_streak = 0
    max_reads = n_expected * 2 + 32
    for _ in range(max_reads):
        sa = net.read("scene", 0)
        if sa.get("eof"):
            eof_streak += 1
            # stragglers can still arrive after an eof; only stop once we've either
            # got everything or seen several eofs in a row (truly drained).
            if len(annotations) >= n_expected or eof_streak >= 8:
                break
            continue
        eof_streak = 0
        annotations.append(sa)
        if len(annotations) >= n_expected:
            break

    if raw_out:
        with open(raw_out, "w") as f:
            for sa in annotations:
                f.write(json.dumps(sa) + "\n")

    result = build_minimal(annotations, movie, fps, cfg)
    if stop:
        net.stop()
    return result


def build_minimal(annotations, movie, fps, cfg) -> dict:
    """Reproduce result_sink's minimal schema: per-actor annealed [start,end] windows.

    Mirrors ResultSinkFunc::build_actor_windows — merge each actor's detection
    timestamps into windows, bridging gaps shorter than anneal_sec.
    """
    anneal = float(cfg.get("anneal_sec", 10.0))
    info = {}           # actor_idx -> identity fields
    times = {}          # actor_idx -> [timestamps]
    for sa in annotations:
        for a in sa["visible_actors"]:
            if a["actor_idx"] < 0:
                continue
            info[a["actor_idx"]] = a
            times.setdefault(a["actor_idx"], []).append(sa["timestamp_sec"])

    actors = []
    for idx, ts in times.items():
        ts.sort()
        scenes = []
        ws = we = ts[0]
        for t in ts[1:]:
            if t - we > anneal:
                scenes.append([ws, we])
                ws = t
            we = t
        scenes.append([ws, we])
        a = info[idx]
        actors.append({
            "name": a["name"], "imdb_id": a["imdb_id"], "tmdb_id": a["tmdb_id"],
            "jellyfin_id": a["jellyfin_id"], "scenes": scenes,
        })

    return {"schema_version": 1, "movie": movie, "sample_fps": fps,
            "anneal_sec": anneal, "actors": actors}


CFG_KEYS = ["detector_conf", "prob_threshold", "match_prior", "match_threshold", "match_ratio",
            "match_ratio_ceil", "track_alpha", "track_min_iou", "track_max_embed_dist",
            "track_max_frames_missing", "cut_revive_sim", "cut_inactive_max_frames",
            "extinction_sec", "anneal_sec"]


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")
    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
    # 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, stop=True,
                    raw_out=args.raw_out)
    Path(args.out).write_text(json.dumps(result, indent=2))
    print(f"[replay] {len(result['actors'])} actors → {args.out}", file=sys.stderr)


if __name__ == "__main__":
    main()