scene-actor-extraction/scripts/scene_detector/train_scene_boundary.py

#!/usr/bin/env python3
"""
train_scene_boundary.py — learn a scene-boundary detector from per-frame RGB
histograms (video tower) and per-second audio log-PSD (audio tower), against
Amazon X-Ray scene boundaries.

Motivation: the shipped grayscale histogram-correlation cut detector is blind on
low-contrast grades — on Scarface it fired ONCE in 10,204 frames, so flood-fill
presence (which snaps to detected boundaries) floods every actor across the whole
film (P=26%). X-Ray ships real scene boundaries (scenes.csv); the dumps carry a
per-frame RGB histogram (frames/rgb_hist), and extract_audio_features.py provides
a per-second audio log-PSD. This learns a per-second boundary probability.

TWO-TOWER, ABLATABLE. We do NOT assume audio helps video — we measure it. Each
modality has its own encoder+BiLSTM; --modality selects video / audio / fused
(both towers concatenated before a shared head). The script reports all three
arms on the held-out films so the ablation decides whether audio supports video.

Video features per second: rgb_hist (96) + L1 deltas to t-1,t-2,t+1 + per-channel
correlation to t-1. Audio features: the log-PSD row (+ its L1 delta to t-1).
Label: 1 if an X-Ray scene starts within ±TOL_SEC of t.

Usage:
  python scripts/scene_detector/train_scene_boundary.py \
      --manifest experiments/manifests/films_LVFace_opencv5.json \
      --audio-dir experiments/dumps/audio_features \
      --holdout Scarface Sound_of_Metal \
      --modality all --out experiments/results/scene_boundary
"""
from __future__ import annotations
import argparse, csv, json, sys
from pathlib import Path

import h5py
import numpy as np
import torch
import torch.nn as nn

TOL_SEC = 2.0
BINS = 32          # per channel, matches embedding_dump_node.hpp kHistBins
RAMP_SCALES = [2, 4, 6, 8, 10]   # multi-scale matched-filter half-widths (seconds)
SCENE_TAU = 205.0   # corpus mean X-Ray scene length (central-60min); debounce scale


def debounce_phase(delta_signal: np.ndarray, tau: float = SCENE_TAU,
                   peak_pct: float = 90.0) -> np.ndarray:
    """A scene-length-scaled 'how overdue is a boundary' feature, [T,2].

    Encodes the prior that scenes don't restart moments apart. From the strong
    peaks of a change signal (the presumed boundaries so far), track time since
    the last peak and turn it into:
      phase = min(1, dt/tau)   — 0 just after a boundary (suppress), 1 when a new
                                 one is overdue (permit), rising over ~one mean
                                 scene length (tau).
      decay = exp(-dt/tau)     — the complementary refractory (high right after,
                                 decaying away). Two views of the same clock so
                                 the LSTM can use whichever helps.
    Reference peaks come from the change signal itself (not the model's own
    output), so the feature is static and causal-ish (uses only |Δ| already in
    the sequence)."""
    T = len(delta_signal)
    thr = np.percentile(delta_signal, peak_pct)
    # Vectorised time-since-last-peak: index of the most recent peak at or before
    # each t (running max of peak indices), then dt = t - that index.
    idx = np.arange(T)
    peak_idx = np.where(delta_signal > thr, idx, -1)
    last = np.maximum.accumulate(peak_idx)          # most recent peak index ≤ t
    dt = (idx - last).astype(np.float32)
    dt[last < 0] = tau                              # before the first peak: treat as "overdue"
    phase = np.minimum(1.0, dt / tau)
    decay = np.exp(-dt / tau)
    return np.stack([phase, decay], 1).astype(np.float32)


def ramp_bank(series: np.ndarray) -> np.ndarray:
    """Antisymmetric matched-filter responses at RAMP_SCALES → [T, len(scales)].

    A scene boundary is a step in the feature series; a signed ramp kernel
    convolved with it responds at the transition and ~0 inside a stable scene.
    Different films' boundaries peak at different scales (measured: sharp cuts at
    H=2s, gradual shifts wider), so we hand the model the whole bank and let it
    weight the scales rather than committing to one width."""
    # Vectorised: the ramp response at t is || sum_l w(l)·series[t+l] ||, i.e. a
    # 1D correlation of the kernel with each feature bin, then an L2 over bins. Do
    # it as one convolution per bin (np.convolve, 'same') instead of the per-frame
    # Python loop — ~100x faster, which matters at ~60k frames × 9 films.
    T, D = series.shape
    out = np.zeros((T, len(RAMP_SCALES)), np.float32)
    for k, H in enumerate(RAMP_SCALES):
        lags = np.arange(-H, H + 1)
        w = (np.sign(lags) * (np.abs(lags) / max(H, 1))).astype(np.float64)
        # correlation = convolution with the reversed kernel; ramp is antisym so
        # reversing negates it — sign folds into the L2 norm, so either is fine.
        acc = np.zeros((T, D))
        for d in range(D):
            acc[:, d] = np.convolve(series[:, d], w[::-1], mode="same")
        out[:, k] = np.linalg.norm(acc, axis=1)
    return out


# ── data ──────────────────────────────────────────────────────────────────────

def load_xray_boundaries(xray_dir: str) -> list[float]:
    starts = []
    with open(Path(xray_dir) / "scenes.csv", newline="") as f:
        for r in csv.DictReader(f):
            s = float(r["start"]) / 1000.0
            if s > 0.5:
                starts.append(s)
    return sorted(starts)


def _znorm(s):
    return (s - s.mean(0)) / (s.std(0) + 1e-6)


def video_features(hist: np.ndarray) -> np.ndarray:
    """DELTA-FORWARD video features.

    Measured on the corpus: the raw 96-bin histogram barely separates X-Ray
    boundaries (~1.4x boundary response) — it encodes what the frame *looks like*,
    not that it *changed* — while the symmetric histogram delta |hist(t+k)-hist(t-k)|
    separates them strongly (|Δ 1s| ~4-5x). Feeding 96 dims of raw content
    diluted the LSTM, so we drop it and lead with multi-scale symmetric deltas,
    keeping only a compact per-channel-energy summary as context.

    Channels:
      - symmetric L1 delta |hist(t+k) - hist(t-k)| at k=1,2,4,8s (the boundary cue)
      - per-channel correlation to the previous second (3)
      - the multi-scale antisymmetric ramp bank (regional step response)
      - 3-D per-channel total energy (compact content context, not the full hist)
    """
    T = hist.shape[0]
    def sym_delta(k):
        fwd = np.roll(hist, -k, 0); fwd[-k:] = hist[-1]
        bwd = np.roll(hist,  k, 0); bwd[:k] = hist[0]
        return np.abs(fwd - bwd).sum(1, keepdims=True)
    deltas = np.concatenate([sym_delta(k) for k in (1, 2, 4, 8)], 1)
    p1 = np.roll(hist, 1, 0); p1[0] = hist[0]
    corr = np.zeros((T, 3), np.float32)
    for c in range(3):
        a = hist[:, c*BINS:(c+1)*BINS]; b = p1[:, c*BINS:(c+1)*BINS]
        am, bm = a - a.mean(1, keepdims=True), b - b.mean(1, keepdims=True)
        corr[:, c] = (am*bm).sum(1) / (np.sqrt((am*am).sum(1)*(bm*bm).sum(1))+1e-9)
    energy = np.stack([hist[:, c*BINS:(c+1)*BINS].sum(1) for c in range(3)], 1)
    # scene-length-scaled debounce: 'how overdue is a boundary', from the |Δ1s|
    # change signal. Encodes that scenes don't restart moments apart (tau=205s).
    debounce = debounce_phase(deltas[:, 0])
    return np.concatenate([deltas, corr, ramp_bank(_znorm(hist)), energy, debounce],
                          1).astype(np.float32)


def audio_features(psd: np.ndarray) -> np.ndarray:
    """DELTA-FORWARD audio features (same principle as video).

    The raw log-PSD is spectral CONTENT (what the audio sounds like), which the DE
    cutter showed barely localizes X-Ray boundaries. Lead with the CHANGE in the
    spectrum — symmetric PSD deltas |psd(t+k)-psd(t-k)| at several scales — plus
    the ramp bank and a compact total-energy summary; drop the full raw PSD.
    """
    def sym_delta(k):
        fwd = np.roll(psd, -k, 0); fwd[-k:] = psd[-1]
        bwd = np.roll(psd,  k, 0); bwd[:k] = psd[0]
        return np.abs(fwd - bwd).sum(1, keepdims=True)
    deltas = np.concatenate([sym_delta(k) for k in (1, 2, 4, 8)], 1)
    energy = psd.sum(1, keepdims=True)
    debounce = debounce_phase(deltas[:, 0])
    return np.concatenate([deltas, ramp_bank(_znorm(psd)), energy, debounce],
                          1).astype(np.float32)


def build_film(dump: str, xray_dir: str, audio_dir: str | None):
    with h5py.File(dump, "r") as f:
        if "frames/rgb_hist" not in f:
            raise SystemExit(f"{dump}: no frames/rgb_hist — re-dump with the "
                             f"RGB-histogram build of dump_embeddings.")
        hist = f["frames/rgb_hist"][:].astype(np.float32)
        ts = f["frames/timestamp_sec"][:]
        is_cut = f["frames/is_cut"][:].astype(np.int64)
    V = video_features(hist)
    A = None
    if audio_dir:
        slug = Path(dump).stem.replace("dump_", "")
        ap = Path(audio_dir) / f"{slug}.npz"
        if ap.exists():
            z = np.load(ap); af = z["feat"]
            # align audio (per-second) to the video frame grid by index; pad/truncate
            T = len(ts); B = af.shape[1]
            aligned = np.zeros((T, B), np.float32)
            m = min(T, len(af)); aligned[:m] = af[:m]
            A = audio_features(aligned)
    y = np.zeros(len(ts), np.float32)
    for b in load_xray_boundaries(xray_dir):
        y[np.abs(ts - b) <= TOL_SEC] = 1.0
    return V, A, y, is_cut, ts


# ── model ─────────────────────────────────────────────────────────────────────

class Tower(nn.Module):
    """Per-second encoder → BiLSTM → per-timestep embedding."""
    def __init__(self, in_dim, hidden=64, out=64):
        super().__init__()
        self.enc = nn.Sequential(nn.Linear(in_dim, hidden), nn.ReLU())
        self.lstm = nn.LSTM(hidden, out, batch_first=True, bidirectional=True)
    def forward(self, x):
        h, _ = self.lstm(self.enc(x))
        return h                       # [B,T,2*out]


class BoundaryNet(nn.Module):
    def __init__(self, v_dim, a_dim, modality):
        super().__init__()
        self.modality = modality
        feat = 0
        if modality in ("video", "fused"):
            self.vtower = Tower(v_dim); feat += 128
        if modality in ("audio", "fused"):
            self.atower = Tower(a_dim); feat += 128
        self.head = nn.Sequential(nn.Linear(feat, 32), nn.ReLU(), nn.Linear(32, 1))
    def forward(self, v, a):
        parts = []
        if self.modality in ("video", "fused"): parts.append(self.vtower(v))
        if self.modality in ("audio", "fused"): parts.append(self.atower(a))
        return self.head(torch.cat(parts, -1)).squeeze(-1)


def nms_peaks(prob, thr=0.5, min_gap=5):
    """Collapse each run of adjacent above-threshold seconds to its single peak.
    Without this, a model that fires 5 consecutive seconds around one true
    boundary is scored as 1 TP + 4 FP — an aggregation artifact, not an error."""
    cand = np.where(prob > thr)[0]
    if len(cand) == 0:
        return []
    peaks, group = [], [cand[0]]
    for c in cand[1:]:
        if c - group[-1] <= min_gap:
            group.append(c)
        else:
            peaks.append(group[int(np.argmax(prob[group]))]); group = [c]
    peaks.append(group[int(np.argmax(prob[group]))])
    return peaks


def prf(prob_or_pred, y, tol=2, thr=0.5):
    """Boundary P/R/F1 with NMS peak aggregation. Accepts a probability series
    (model output) or a 0/1 array (is_cut baseline); NMS collapses each run of
    above-threshold seconds to one peak either way."""
    P = np.array(nms_peaks(np.asarray(prob_or_pred, float), thr=thr))
    T = np.where(y > 0.5)[0]
    if len(P) == 0 or len(T) == 0: return 0., 0., 0.
    tp_p = sum(any(abs(p-t) <= tol for t in T) for p in P)
    tp_t = sum(any(abs(p-t) <= tol for p in P) for t in T)
    pr, rc = tp_p/len(P), tp_t/len(T)
    return pr, rc, (2*pr*rc/(pr+rc) if pr+rc else 0.)


def train_arm(modality, tr, va, v_dim, a_dim, vmu, vsd, amu, asd, epochs, dev):
    model = BoundaryNet(v_dim, a_dim, modality).to(dev)
    opt = torch.optim.Adam(model.parameters(), lr=1e-3, weight_decay=1e-5)
    pos = sum((y > .5).sum() for *_, y, _, _ in tr)
    neg = sum((y <= .5).sum() for *_, y, _, _ in tr)
    lossf = nn.BCEWithLogitsLoss(pos_weight=torch.tensor([neg/max(pos,1)], device=dev))
    def vt(V): return torch.tensor((V-vmu)/vsd, dtype=torch.float32, device=dev).unsqueeze(0)
    def at(A): return torch.tensor((A-amu)/asd, dtype=torch.float32, device=dev).unsqueeze(0)
    for ep in range(epochs):
        model.train()
        for V, A, y, _, _ in tr:
            opt.zero_grad()
            logit = model(vt(V), at(A) if A is not None else None)
            loss = lossf(logit, torch.tensor(y, device=dev).unsqueeze(0))
            loss.backward(); opt.step()
    model.eval(); rows = {}
    with torch.no_grad():
        for slug, V, A, y, is_cut, ts in va:
            prob = torch.sigmoid(model(vt(V), at(A) if A is not None else None))[0].cpu().numpy()
            rows[slug] = prf(prob, y)   # raw prob → NMS picks peaks by height
    return model, rows


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--manifest", required=True)
    ap.add_argument("--audio-dir", default="experiments/dumps/audio_features")
    ap.add_argument("--holdout", nargs="+", default=["Scarface", "Sound_of_Metal"])
    ap.add_argument("--modality", choices=["video","audio","fused","all"], default="all")
    ap.add_argument("--out", default="experiments/results/scene_boundary")
    ap.add_argument("--epochs", type=int, default=250)
    ap.add_argument("--seed", type=int, default=0)
    args = ap.parse_args()
    torch.manual_seed(args.seed); np.random.seed(args.seed)

    films = json.load(open(args.manifest))
    def load(rows):
        out = []
        for f in rows:
            V, A, y, is_cut, ts = build_film(f["dump"], f["xray"], args.audio_dir)
            out.append((f["slug"], V, A, y, is_cut, ts))
        return out
    tr = load([f for f in films if f["slug"] not in args.holdout])
    va = load([f for f in films if f["slug"] in args.holdout])
    has_audio = all(t[2] is not None for t in tr+va)
    print(f"[scene] train {len(tr)} / holdout {args.holdout}; audio={'yes' if has_audio else 'MISSING'}",
          file=sys.stderr)

    allV = np.concatenate([t[1] for t in tr], 0)
    vmu, vsd = allV.mean(0), allV.std(0)+1e-6; v_dim = allV.shape[1]
    if has_audio:
        allA = np.concatenate([t[2] for t in tr], 0)
        amu, asd = allA.mean(0), allA.std(0)+1e-6; a_dim = allA.shape[1]
    else:
        amu = asd = None; a_dim = 1

    # strip index tuples for train_arm (expects V,A,y,is_cut,ts)
    trA = [(t[1],t[2],t[3],t[4],t[5]) for t in tr]
    dev = "cuda" if torch.cuda.is_available() else "cpu"
    modes = ["video","audio","fused"] if args.modality=="all" else [args.modality]
    if not has_audio: modes = [m for m in modes if m == "video"] or ["video"]

    # grayscale-0.70 baseline (is_cut) on holdout
    print("\n=== held-out scene-boundary detection (P/R/F1, ±2s) ===")
    print(f"{'film':26s} " + "  ".join(f"{m:>16s}" for m in modes) + f"  {'grayscale-0.70':>16s}")
    Path(args.out).mkdir(parents=True, exist_ok=True)
    results = {m: train_arm(m, trA, va, v_dim, a_dim, vmu, vsd, amu, asd, args.epochs, dev)
               for m in modes}
    report = {"holdout": args.holdout, "tol_sec": TOL_SEC, "modalities": {}, "films": {}}
    for slug, V, A, y, is_cut, ts in va:
        cells = []
        for m in modes:
            p,r,f = results[m][1][slug]
            cells.append(f"{p*100:4.0f}/{r*100:4.0f}/{f*100:4.0f}")
            report["films"].setdefault(slug, {})[m] = {"P":p,"R":r,"F1":f}
        bp,br,bf = prf(is_cut, y)
        report["films"].setdefault(slug, {})["grayscale"] = {"P":bp,"R":br,"F1":bf}
        print(f"{slug:26s} " + "  ".join(f"{c:>16s}" for c in cells) +
              f"  {bp*100:4.0f}/{br*100:4.0f}/{bf*100:4.0f}")
    # macro-mean F1 per modality across holdout
    print("\nmacro-mean holdout F1:")
    for m in modes:
        mf = np.mean([results[m][1][s][2] for s,*_ in va])
        report["modalities"][m] = float(mf)
        print(f"  {m:8s} {mf*100:.1f}%")
    bf = np.mean([prf(t[4], t[3])[2] for t in va])
    report["modalities"]["grayscale"] = float(bf)
    print(f"  {'grayscale':8s} {bf*100:.1f}%")
    # save the best arm
    best = max(modes, key=lambda m: report["modalities"][m])
    torch.save({"state": results[best][0].state_dict(), "modality": best,
                "vmu":vmu,"vsd":vsd,"amu":amu,"asd":asd,"v_dim":v_dim,"a_dim":a_dim},
               Path(args.out)/"boundary_net.pt")
    json.dump(report, open(Path(args.out)/"report.json","w"), indent=2)
    print(f"\n[scene] best={best}; model+report → {args.out}", file=sys.stderr)


if __name__ == "__main__":
    main()