experiment-2107/scripts/temporal_scripts/fs_spot_dense.py

"""Dense-supervision TAS on FS-Spot -- the counterpart to fs_spot_ctc.py.

Why this exists: CTC converged to loss 0.338 on this data and still scored 0.7%
localization recall. It emitted the right elements, in the right ORDER, all
crammed into the first ~20 frames of a 4,275-frame clip. That is not a bug --
CTC's loss marginalizes over every monotonic alignment, so spikes at frame 0
score identically to spikes at the true frame. With targets at 0.5% density and a
bidirectional encoder that sees the whole clip, nothing pushes back.

Dense per-frame cross-entropy supplies exactly the pressure CTC lacks: predicting
Axel at frame 0 when the Axel is at frame 578 is directly penalized. Same encoder
and same features as the CTC run, so any difference is the loss, not the model.

Two things this must handle that CTC did not:
  * 48:1 imbalance (97.94% background). Unweighted CE predicts all-background and
    reports 97.9% "accuracy" while finding nothing. Hence class weights.
  * Fragmented predictions. Frame-wise CE has no notion of a contiguous segment,
    so it produces speckle. TMSE (the MS-TCN smoothing term) penalizes abrupt
    frame-to-frame log-prob changes and glues segments together.

Run:  python3 -m temporal_scripts.fs_spot_dense [--epochs 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_ctc import STRIDE, CTCTagger, skeleton_features
from temporal_scripts.fs_spot_data import OUT_DIR, SPOT_TAXONOMY

DEVICE = "cuda:0" if torch.cuda.is_available() else "cpu"
NUM_CLASSES = len(SPOT_TAXONOMY)  # 8: NONE + 6 jumps + Spin


def downsample_labels(y: np.ndarray, stride: int = STRIDE) -> np.ndarray:
    """(T,) -> (T//stride,), keeping any element that appears in each window.

    Elements are only ~18 frames (~4 windows at stride 4), so plain y[::stride]
    can drop one entirely by landing between samples. Take the max instead: any
    non-background label in the window survives.
    """
    n = len(y) // stride
    return y[: n * stride].reshape(n, stride).max(axis=1)


class FSSpotDense(Dataset):
    def __init__(self, files: list[Path]):
        self.items = []
        for f in files:
            d = np.load(f, allow_pickle=True)
            self.items.append(
                {
                    "feats": skeleton_features(d["skeleton"]),
                    "y": downsample_labels(d["frame_labels"]),
                    "takeoff": d["element_takeoff"],
                    "landing": d["element_landing"],
                    "labels": [str(x) for x in d["element_label"]],
                    "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)
    Tem = T // STRIDE
    F_ = batch[0]["feats"].shape[1]
    x = torch.zeros(len(batch), T, F_)
    y = torch.full((len(batch), Tem), -100, dtype=torch.long)  # -100 = ignored by CE
    for i, b in enumerate(batch):
        t = len(b["feats"])
        x[i, :t] = torch.from_numpy(b["feats"])
        ye = b["y"]
        y[i, : len(ye)] = torch.from_numpy(ye)
    return {"x": x, "y": y, "meta": batch}


def tmse_loss(logp: torch.Tensor, tau: float = 4.0) -> torch.Tensor:
    """Truncated MSE over adjacent-frame log-prob deltas (MS-TCN's smoothing term).

    Clamping at tau keeps a genuine action boundary -- where a large jump is
    correct -- from being penalized into oblivion.
    """
    d = (logp[:, 1:] - logp[:, :-1]) ** 2
    return torch.clamp(d, max=tau**2).mean()


def segments_from_frames(pred: np.ndarray) -> list[tuple[int, int, int]]:
    """Collapse a per-frame prediction into (class, start, end) runs, dropping background."""
    out = []
    cur, start = pred[0], 0
    for t in range(1, len(pred) + 1):
        if t == len(pred) or pred[t] != cur:
            if cur != 0:
                out.append((int(cur), start, t - 1))
            if t < len(pred):
                cur, start = pred[t], t
    return out


def evaluate(model: nn.Module, loader: DataLoader, tol_s: float = 1.0, min_len: int = 2) -> dict:
    """Same localization metric as the CTC run, so the two are directly comparable.

    A hit = a predicted segment whose centre is within tol_s of a true element's
    [takeoff, landing] window, with matching class, one-to-one.
    """
    model.eval()
    hits = gt_total = pred_total = 0
    offsets = []
    frame_correct = frame_total = 0
    with torch.no_grad():
        for b in loader:
            logits = model(b["x"].to(DEVICE))
            pred = logits.argmax(-1).cpu().numpy()
            for i, m in enumerate(b["meta"]):
                n = min(len(m["y"]), pred.shape[1])
                p = pred[i][:n]
                frame_correct += int((p == m["y"][:n]).sum())
                frame_total += n
                segs = [s for s in segments_from_frames(p) if s[2] - s[1] + 1 >= min_len]
                gt = [
                    (lab, int(t0), int(t1))
                    for lab, t0, t1 in zip(m["labels"], m["takeoff"], m["landing"])
                ]
                gt_total += len(gt)
                pred_total += len(segs)
                used = set()
                for cls, s0, s1 in segs:
                    centre = (s0 + s1) / 2 * STRIDE  # emission grid -> input frames
                    for j, (lab, t0, t1) in enumerate(gt):
                        if j in used or SPOT_TAXONOMY[cls] != lab:
                            continue
                        dist = max(t0 - centre, 0, centre - t1) / 25.0
                        if dist <= tol_s:
                            hits += 1
                            used.add(j)
                            offsets.append(abs(centre - (t0 + t1) / 2))
                            break
    return {
        "loc_recall": hits / max(gt_total, 1),
        "loc_prec": hits / max(pred_total, 1),
        "frame_acc": frame_correct / max(frame_total, 1),
        "median_offset_s": float(np.median(offsets) / 25.0) if offsets else float("nan"),
        "n_pred": pred_total,
        "n_gt": gt_total,
    }


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--epochs", type=int, default=60)
    ap.add_argument("--batch", type=int, default=4)
    ap.add_argument("--lr", type=float, default=3e-4)
    ap.add_argument("--lambda-tmse", type=float, default=0.15)
    ap.add_argument("--unidirectional", action="store_true")
    ap.add_argument("--tag", type=str, default="dense")
    args = ap.parse_args()

    files = sorted(Path(p) for p in glob.glob(str(OUT_DIR / "*.npz")))
    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)} | train {len(train_f)} ({len(train_bcs)} bc) | val {len(val_f)} ({val_bcs})")

    tr, va = FSSpotDense(train_f), FSSpotDense(val_f)
    in_features = tr[0]["feats"].shape[1]

    # Inverse-frequency class weights. Background is 48x everything else and Loop
    # is 0.048% of frames; without this the model predicts all-background and
    # reports 97.9% frame accuracy while finding nothing. sqrt tempers the tail so
    # Loop's ~300 frames don't dominate the gradient outright.
    counts = np.zeros(NUM_CLASSES)
    for it in tr.items:
        for k, v in zip(*np.unique(it["y"], return_counts=True)):
            counts[int(k)] += v
    w = 1.0 / np.sqrt(np.maximum(counts, 1))
    w = w / w.sum() * NUM_CLASSES
    print("class weights:", {SPOT_TAXONOMY[i]: round(float(w[i]), 2) for i in range(NUM_CLASSES)})
    weight = torch.tensor(w, dtype=torch.float32, device=DEVICE)

    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, num_classes=NUM_CLASSES,
                      bidirectional=not args.unidirectional).to(DEVICE)
    print(f"  LSTM: {'UNIdirectional' if args.unidirectional else 'BIdirectional'}")
    opt = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=1e-4)
    sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=args.epochs)

    best = -1.0
    for ep in range(1, args.epochs + 1):
        model.train()
        tot = nb = 0
        for b in tl:
            opt.zero_grad()
            logits = model(b["x"].to(DEVICE))
            y = b["y"].to(DEVICE)
            # The strided conv stem's output length follows the conv formula, which
            # is T//STRIDE +/- 1 rather than exactly T/4. Crop both to the common
            # length so frame t of the logits lines up with frame t of the labels.
            n = min(logits.shape[1], y.shape[1])
            logits, y = logits[:, :n], y[:, :n]
            ce = F.cross_entropy(logits.reshape(-1, NUM_CLASSES), y.reshape(-1),
                                 weight=weight, ignore_index=-100)
            logp = F.log_softmax(logits, dim=-1)
            loss = ce + args.lambda_tmse * tmse_loss(logp)
            if not torch.isfinite(loss):
                continue
            loss.backward()
            nn.utils.clip_grad_norm_(model.parameters(), 5.0)
            opt.step()
            tot += float(loss.detach()); 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} | frame_acc {m['frame_acc']*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 | pred {m['n_pred']}/gt {m['n_gt']}",
                flush=True,
            )
            if m["loc_recall"] > best:
                best = m["loc_recall"]
                torch.save({"state_dict": model.state_dict(), "in_features": in_features,
                            "bidirectional": not args.unidirectional},
                           OUT_DIR.parent / f"{args.tag}_model.pt")
    print(f"\nbest val loc_recall {best*100:.1f}% -> {OUT_DIR.parent/(args.tag+'_model.pt')}")
    return 0


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