experiment-2107/scripts/temporal_scripts/fs_spot_ctc.py

"""CTC training on FS-Spot: learn to timestamp skating elements from transcripts.

The point of CTC here: it needs only the ORDERED TRANSCRIPT per clip
(["Axel", "Flip", "Toeloop", ...]) -- no frame numbers -- and recovers the
alignment itself by marginalizing over every way the transcript could map onto
the frames. Timestamps fall out of the emissions at inference.

Background is NOT a class. CTC's blank token is what covers the ~98% of frames
between elements; adding an explicit NONE class alongside blank would be
double-counting. Classes are just the 7 element types.

We DO have ground-truth boundaries here (element_takeoff/element_landing), which
is what makes this worth running: CTC never sees them, so they are a clean test
of whether the recovered timestamps are real. That is the experiment.

    features   (T, F)  per-frame, scale/translation-normalized skeleton + motion
    target     (L,)    ordered element ids, 1..7  (0 is reserved for CTC blank)
    emissions  (T/4, C+1)  encoder stride-4 downsamples; every input frame is
                           still consumed, the emission grid is just coarser
                           (160ms), which is fine against ~700ms jumps.

Run:  python3 -m temporal_scripts.fs_spot_ctc [--epochs N] [--limit N]
"""

from __future__ import annotations

import argparse
import glob
import sys
from pathlib import Path

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader, Dataset

from temporal_scripts.fs_spot_data import OUT_DIR

DEVICE = "cuda:0" if torch.cuda.is_available() else "cpu"

# CTC blank is 0, so element ids start at 1.
CTC_CLASSES = ["Axel", "Flip", "Loop", "Lutz", "Salchow", "Toeloop", "Spin"]
CTC_LABEL_TO_ID = {name: i + 1 for i, name in enumerate(CTC_CLASSES)}
CTC_ID_TO_LABEL = {i + 1: name for i, name in enumerate(CTC_CLASSES)}
NUM_CTC_CLASSES = len(CTC_CLASSES) + 1  # + blank

STRIDE = 4  # encoder temporal downsample

# COCO-17 indices
L_SHO, R_SHO, L_HIP, R_HIP = 5, 6, 11, 12


def skeleton_features(skel: np.ndarray) -> np.ndarray:
    """(T,17,3) normalized-pixel keypoints -> (T, F) scale/translation invariant features.

    Broadcast cameras pan and zoom constantly, so raw normalized pixel coords
    encode camera work as much as skating. Center each frame on the hip midpoint
    and divide by torso length, which cancels both. Motion is then added back as
    explicit velocity channels -- of the normalized coords (body-relative motion,
    e.g. limbs) and of the hip centre itself (whole-body translation, which is
    what actually spikes on a jump).
    """
    xy = skel[:, :, :2].astype(np.float32)
    conf = skel[:, :, 2].astype(np.float32)

    hip = (xy[:, L_HIP] + xy[:, R_HIP]) / 2.0  # (T,2)
    sho = (xy[:, L_SHO] + xy[:, R_SHO]) / 2.0
    torso = np.linalg.norm(sho - hip, axis=1, keepdims=True)  # (T,1)
    # Frames with no detection are all-zero -> torso 0. Guard, then let the
    # invalid mask tell the model which frames are junk rather than hiding it.
    valid = (conf > 0).any(axis=1) & (torso[:, 0] > 1e-5)
    torso = np.maximum(torso, 1e-3)

    centered = (xy - hip[:, None, :]) / torso[:, None, :]  # (T,17,2)
    centered[~valid] = 0.0

    def d(a: np.ndarray) -> np.ndarray:
        return np.concatenate([np.zeros((1,) + a.shape[1:], a.dtype), np.diff(a, axis=0)], axis=0)

    vel = d(centered)  # body-relative joint motion
    hip_v = d(hip)  # whole-body translation (the jump signal)
    hip_a = d(hip_v)

    feats = np.concatenate(
        [
            centered.reshape(len(xy), -1),  # 34
            vel.reshape(len(xy), -1),       # 34
            conf,                            # 17
            hip,                             # 2  (where in frame -- weak, but free)
            hip_v,                           # 2
            hip_a,                           # 2
            torso,                           # 1  (scale = camera zoom proxy)
            valid[:, None].astype(np.float32),  # 1  (explicit "no detection here")
        ],
        axis=1,
    )
    return np.nan_to_num(feats, nan=0.0, posinf=0.0, neginf=0.0).astype(np.float32)


class FSSpotCTC(Dataset):
    def __init__(self, files: list[Path]):
        self.items = []
        for f in files:
            d = np.load(f, allow_pickle=True)
            labels = [str(x) for x in d["element_label"]]
            target = [CTC_LABEL_TO_ID[x] for x in labels if x in CTC_LABEL_TO_ID]
            if not target:
                continue
            self.items.append(
                {
                    "file": f,
                    "feats": skeleton_features(d["skeleton"]),
                    "target": np.array(target, dtype=np.int64),
                    "takeoff": d["element_takeoff"],
                    "landing": d["element_landing"],
                    "labels": labels,
                    "name": f.stem,
                }
            )

    def __len__(self) -> int:
        return len(self.items)

    def __getitem__(self, i: int) -> dict:
        return self.items[i]


def collate(batch: list[dict]) -> dict:
    T = max(len(b["feats"]) for b in batch)
    F_ = batch[0]["feats"].shape[1]
    x = torch.zeros(len(batch), T, F_)
    in_len, tg_len, targets = [], [], []
    for i, b in enumerate(batch):
        t = len(b["feats"])
        x[i, :t] = torch.from_numpy(b["feats"])
        in_len.append(t // STRIDE)  # emission length after the strided stem
        tg_len.append(len(b["target"]))
        targets.append(torch.from_numpy(b["target"]))
    return {
        "x": x,
        "targets": torch.cat(targets),
        "input_lengths": torch.tensor(in_len, dtype=torch.long),
        "target_lengths": torch.tensor(tg_len, dtype=torch.long),
        "meta": batch,
    }


class CTCTagger(nn.Module):
    """Encoder (conv stem + LSTM) -> per-frame class logits.

    The encoder is BOTH parts: the conv stem does local feature extraction and the
    stride-4 downsample, the LSTM adds temporal context. Only the LSTM's direction
    is configurable, because only the LSTM has unbounded context.

    Stem receptive field: conv(k5,s2) -> 5, conv(k5,s2) -> 13, conv(k3,s1) -> 21
    input frames, centered, i.e. +/-10 frames (~0.4s) of lookahead. That window is
    ~one jump wide (a jump is ~18 frames) and is far too small to let the model
    know at frame 0 what happens at frame 970.

    bidirectional=False is the CTC alignment experiment. CTC's loss marginalizes
    over alignments, so it cannot distinguish emitting the transcript at frame 0
    from emitting it at the truth -- both collapse to the same label sequence and
    score the same. A BiLSTM's backward pass has read all 4,275 frames by step 0,
    which makes the frame-0 dump representable AND cheaper to learn than genuine
    onset detection. Going unidirectional removes the global lookahead, so firing
    early would require predicting the future. If the alignment-invariance
    diagnosis is right, this should localize; if it stays ~1%, the diagnosis is
    wrong. Expect emissions biased LATE regardless -- streaming CTC must see
    enough of an event to commit to a class.
    """

    def __init__(self, in_features: int, num_classes: int = NUM_CTC_CLASSES, hidden: int = 256,
                 bidirectional: bool = True):
        super().__init__()
        self.stem = nn.Sequential(
            nn.Conv1d(in_features, 128, 5, stride=2, padding=2), nn.BatchNorm1d(128), nn.GELU(),
            nn.Conv1d(128, 256, 5, stride=2, padding=2), nn.BatchNorm1d(256), nn.GELU(),
            nn.Conv1d(256, 256, 3, padding=1), nn.BatchNorm1d(256), nn.GELU(),
        )
        self.bidirectional = bidirectional
        self.rnn = nn.LSTM(256, hidden, num_layers=2, batch_first=True,
                           bidirectional=bidirectional, dropout=0.2)
        self.out = nn.Linear(hidden * (2 if bidirectional else 1), num_classes)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.stem(x.transpose(1, 2)).transpose(1, 2)
        x, _ = self.rnn(x)
        return self.out(x)  # (B, T/4, C+1)


def greedy_decode(logp: torch.Tensor, in_len: int) -> list[tuple[int, int]]:
    """Collapse CTC emissions to (class_id, emission_frame) pairs."""
    ids = logp[:in_len].argmax(dim=-1).cpu().numpy()
    out, prev = [], 0
    for t, k in enumerate(ids):
        if k != prev and k != 0:
            out.append((int(k), t))
        prev = k
    return out


def edit_distance(a: list, b: list) -> int:
    prev = list(range(len(b) + 1))
    for i, x in enumerate(a, 1):
        cur = [i]
        for j, y in enumerate(b, 1):
            cur.append(min(prev[j] + 1, cur[j - 1] + 1, prev[j - 1] + (x != y)))
        prev = cur
    return prev[-1]


def evaluate(model: nn.Module, loader: DataLoader, tol_s: float = 1.0) -> dict:
    """Transcript accuracy + whether the emissions land on the true elements.

    CTC never sees a frame number in training, so localization here is emergent.
    A hit = decoded spike within `tol_s` of the true [takeoff, landing] window
    AND the right class.
    """
    model.eval()
    tot_ed = tot_len = exact = n = 0
    hits = gt_total = pred_total = 0
    offsets = []
    with torch.no_grad():
        for b in loader:
            logp = F.log_softmax(model(b["x"].to(DEVICE)), dim=-1)
            for i, m in enumerate(b["meta"]):
                dec = greedy_decode(logp[i], int(b["input_lengths"][i]))
                pred = [CTC_ID_TO_LABEL[k] for k, _ in dec]
                gt = [x for x in m["labels"] if x in CTC_LABEL_TO_ID]
                tot_ed += edit_distance(pred, gt)
                tot_len += len(gt)
                exact += int(pred == gt)
                n += 1

                gt_total += len(gt)
                pred_total += len(pred)
                used = set()
                for k, t_em in dec:
                    t_frame = t_em * STRIDE  # emission grid -> input frames
                    for j, (lab, t0, t1) in enumerate(zip(m["labels"], m["takeoff"], m["landing"])):
                        if j in used or lab != CTC_ID_TO_LABEL[k]:
                            continue
                        # distance to the true window (0 if inside it)
                        dist = max(t0 - t_frame, 0, t_frame - t1) / 25.0
                        if dist <= tol_s:
                            hits += 1
                            used.add(j)
                            offsets.append(t_frame - (t0 + t1) / 2)
                            break
    ser = tot_ed / max(tot_len, 1)
    return {
        "transcript_err": ser,                     # like WER: edits / true elements
        "exact_clip": exact / max(n, 1),
        "loc_recall": hits / max(gt_total, 1),
        "loc_prec": hits / max(pred_total, 1),
        "median_offset_s": float(np.median(np.abs(offsets)) / 25.0) if offsets else float("nan"),
        "n_clips": n,
    }


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--epochs", type=int, default=60)
    ap.add_argument("--limit", type=int, default=None)
    ap.add_argument("--batch", type=int, default=4)
    ap.add_argument("--lr", type=float, default=3e-4)
    ap.add_argument("--unidirectional", action="store_true",
                    help="causal-ish LSTM: removes the global lookahead that lets CTC dump at frame 0")
    ap.add_argument("--tag", type=str, default="ctc", help="checkpoint name")
    args = ap.parse_args()

    files = sorted(Path(p) for p in glob.glob(str(OUT_DIR / "*.npz")))
    if args.limit:
        files = files[: args.limit]
    if len(files) < 10:
        print(f"only {len(files)} clips extracted -- wait for extraction")
        return 1

    # Split BY BROADCAST: clips from one broadcast share camera work, lighting and
    # commentary bumpers, so a random split would leak that and flatter the score.
    by_bc: dict[str, list[Path]] = {}
    for f in files:
        by_bc.setdefault("_".join(f.stem.split("_")[:-3]), []).append(f)
    bcs = sorted(by_bc)
    n_val = max(1, len(bcs) // 4)
    val_bcs, train_bcs = bcs[:n_val], bcs[n_val:]
    train_f = [f for b in train_bcs for f in by_bc[b]]
    val_f = [f for b in val_bcs for f in by_bc[b]]

    print(f"clips: {len(files)} | broadcasts: {len(bcs)}")
    print(f"  train {len(train_f)} clips from {train_bcs}")
    print(f"  val   {len(val_f)} clips from {val_bcs}")

    tr, va = FSSpotCTC(train_f), FSSpotCTC(val_f)
    if not len(tr) or not len(va):
        print("empty split")
        return 1
    in_features = tr[0]["feats"].shape[1]
    print(f"  in_features={in_features} | mean T={np.mean([len(i['feats']) for i in tr.items]):.0f}")

    tl = DataLoader(tr, batch_size=args.batch, shuffle=True, collate_fn=collate)
    vl = DataLoader(va, batch_size=args.batch, shuffle=False, collate_fn=collate)

    model = CTCTagger(in_features, bidirectional=not args.unidirectional).to(DEVICE)
    print(f"  LSTM: {'UNIdirectional (no global lookahead)' if args.unidirectional else 'BIdirectional'}"
          f" | stem lookahead +/-10 frames (0.4s)")
    opt = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=1e-4)
    sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=args.epochs)
    # zero_infinity: a clip whose emission length falls under its target length
    # yields inf loss; drop those grads instead of poisoning the whole batch.
    ctc = nn.CTCLoss(blank=0, zero_infinity=True)

    best = float("inf")
    for ep in range(1, args.epochs + 1):
        model.train()
        tot = nb = 0
        for b in tl:
            opt.zero_grad()
            logp = F.log_softmax(model(b["x"].to(DEVICE)), dim=-1).transpose(0, 1)  # (T,B,C)
            loss = ctc(logp, b["targets"].to(DEVICE), b["input_lengths"], b["target_lengths"])
            if not torch.isfinite(loss):
                continue
            loss.backward()
            nn.utils.clip_grad_norm_(model.parameters(), 5.0)
            opt.step()
            tot += float(loss); nb += 1
        sched.step()
        if ep % 5 == 0 or ep == args.epochs:
            m = evaluate(model, vl)
            print(
                f"ep {ep:3d} | loss {tot/max(nb,1):6.3f} | "
                f"transcript_err {m['transcript_err']:.3f} | exact {m['exact_clip']*100:5.1f}% | "
                f"loc R {m['loc_recall']*100:5.1f}% P {m['loc_prec']*100:5.1f}% | "
                f"med_off {m['median_offset_s']:.2f}s",
                flush=True,
            )
            if tot / max(nb, 1) < best:
                best = tot / max(nb, 1)
                torch.save({"state_dict": model.state_dict(), "in_features": in_features,
                            "bidirectional": not args.unidirectional},
                           OUT_DIR.parent / f"{args.tag}_model.pt")
    print("\nsaved ->", OUT_DIR.parent / f"{args.tag}_model.pt")
    return 0


if __name__ == "__main__":
    sys.exit(main())