experiment-2107/scripts/temporal_scripts/fs_spot_extract.py

"""Extract (T, 17, 3) skeletons + dense frame labels for the 371 FS-Spot clips.

Streams each clip's frames out of its parent broadcast, runs YOLO pose, picks the
skater with a continuity tracker, and writes one npz per clip:

    skeleton     (T, 17, 3)  COCO-17, x/width and y/height normalized, + confidence
    frame_labels (T,)        dense int labels over SPOT_TAXONOMY (0 = NONE)
    box          (T, 4)      the tracked skater's box, xyxy in pixels
    det_count    (T,)        persons detected per frame -- crowd/pan diagnostic

Run:  python3 -m temporal_scripts.fs_spot_extract [--limit N] [--broadcast NAME]

Only broadcasts already present in data/fs_spot/videos are processed, so this can
run alongside the download and be re-run to pick up newly-arrived videos. Clips
whose npz already exists are skipped.
"""

from __future__ import annotations

import argparse
import sys
import time
from pathlib import Path

import numpy as np

from temporal_scripts.fs_spot_data import (
    OUT_DIR,
    VIDEO_DIR,
    Clip,
    available_broadcasts,
    build_index,
    find_video,
    pick_skater,
    stream_window,
)

# ffmpeg pre-scales to this, matching IMGSZ's aspect so YOLO's letterbox is a
# no-op and the frame is resampled once rather than twice. Keypoints are
# normalized by these dims, so [0,1] coords are resolution-independent.
#
# IMGSZ tracks config.py's yolo_imgsz (384). Measured recall ON labelled
# jump/spin frames (the ones that carry the signal), 2 clips, this data:
#     yolo11n @ 384  192 fps   85.98%     <- current: fast PoC setting
#     yolo11m @ 384  191 fps   88.41%     (bigger model is ~free -- GPU isn't saturated)
#     yolo11n @ 640  148 fps   92.07%
#     yolo11m @ 640  130 fps   96.34%
# Resolution dominates model size here. At 384 roughly 1 in 7 element frames has
# no detection at all and lands as a zero row for interpolate_missing to fill --
# concentrated in takeoff/landing motion blur, i.e. exactly the boundaries a TAS
# model must learn. Raise IMGSZ to 640 (and WIDTH/HEIGHT to 640x360) before
# trusting any final number.
WIDTH, HEIGHT = 384, 216
IMGSZ = 384
BATCH = 64


def extract_clip(clip: Clip, model, chunk: int = 128) -> dict:
    """Run pose over one clip and return arrays to save."""
    skel = np.zeros((clip.num_frames, 17, 3), dtype=np.float32)
    boxes = np.zeros((clip.num_frames, 4), dtype=np.float32)
    dets = np.zeros(clip.num_frames, dtype=np.int16)

    video = find_video(clip.broadcast)
    if video is None:
        raise FileNotFoundError(f"no video on disk for broadcast {clip.broadcast}")
    prev_center: np.ndarray | None = None
    t = 0
    for frames in stream_window(
        video, clip.start, clip.num_frames, WIDTH, HEIGHT, chunk=chunk, scale=(WIDTH, HEIGHT)
    ):
        results = []
        for i in range(0, len(frames), BATCH):
            results.extend(
                model.predict(
                    list(frames[i : i + BATCH]),
                    device=0, verbose=False, imgsz=IMGSZ, half=True, batch=BATCH,
                )
            )
        for r in results:
            if t >= clip.num_frames:
                break
            if r.boxes is None or len(r.boxes) == 0 or r.keypoints is None:
                # No person found: leave zeros and drop the track, so the next
                # detection isn't yanked toward a stale centre from before a cut.
                prev_center = None
                t += 1
                continue
            b = r.boxes.xyxy.cpu().numpy()
            cf = r.boxes.conf.cpu().numpy()
            idx = pick_skater(b, cf, prev_center)
            prev_center = np.array([(b[idx, 0] + b[idx, 2]) / 2, (b[idx, 1] + b[idx, 3]) / 2])
            k = r.keypoints.data[idx].cpu().numpy().astype(np.float32)  # (17,3) px x, px y, conf
            k[:, 0] /= WIDTH
            k[:, 1] /= HEIGHT
            skel[t] = k
            boxes[t] = b[idx]
            dets[t] = len(b)
            t += 1

    if t != clip.num_frames:
        raise RuntimeError(f"{clip.name}: decoded {t} frames, expected {clip.num_frames}")

    return {
        "skeleton": skel,
        "frame_labels": clip.frame_labels(),
        "box": boxes,
        "det_count": dets,
        "transcript": np.array(clip.transcript, dtype=object),
        "element_takeoff": np.array([e.takeoff for e in clip.elements], dtype=np.int32),
        "element_landing": np.array([e.landing for e in clip.elements], dtype=np.int32),
        "element_label": np.array([e.label for e in clip.elements], dtype=object),
        "fps": 25.0,
        "broadcast": clip.broadcast,
        "abs_start": clip.start,
    }


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--limit", type=int, default=None, help="process at most N clips")
    ap.add_argument("--broadcast", type=str, default=None, help="only this broadcast")
    ap.add_argument("--weights", type=str, default="yolo11n-pose.pt")
    ap.add_argument("--shard", type=str, default=None, help="i/N -- process only shard i of N")
    args = ap.parse_args()

    from ultralytics import YOLO

    OUT_DIR.mkdir(parents=True, exist_ok=True)
    clips = build_index()

    have = available_broadcasts()
    todo = [
        c for c in clips.values()
        if c.broadcast in have
        and (args.broadcast is None or c.broadcast == args.broadcast)
        and not (OUT_DIR / f"{c.name}.npz").exists()
    ]
    todo.sort(key=lambda c: (c.broadcast, c.start))
    if args.shard:
        i, n = (int(x) for x in args.shard.split("/"))
        todo = [c for j, c in enumerate(todo) if j % n == i]
    if args.limit:
        todo = todo[: args.limit]

    print(f"broadcasts downloaded : {len(have)}/11  {sorted(have)}")
    print(f"clips to process      : {len(todo)}  (done: {len(list(OUT_DIR.glob('*.npz')))}/371)")
    if not todo:
        print("nothing to do")
        return 0

    model = YOLO(args.weights)
    model.to("cuda:0")

    t0 = time.time()
    frames_done = 0
    for i, clip in enumerate(todo, 1):
        ts = time.time()
        try:
            out = extract_clip(clip, model)
        except Exception as exc:
            print(f"[{i}/{len(todo)}] FAIL {clip.name}: {exc}", flush=True)
            continue
        np.savez_compressed(OUT_DIR / f"{clip.name}.npz", **out)
        frames_done += clip.num_frames
        dt = time.time() - ts
        found = float((out["skeleton"][:, :, 2] > 0).any(axis=1).mean())
        elapsed = time.time() - t0
        eta = (sum(c.num_frames for c in todo[i:]) / (frames_done / elapsed)) / 60 if frames_done else 0
        print(
            f"[{i}/{len(todo)}] {clip.name[:52]:52} T={clip.num_frames:5d} "
            f"{clip.num_frames/dt:6.1f} fps  pose_found={found*100:5.1f}%  "
            f"elems={len(clip.elements)}  ETA {eta:5.1f}m",
            flush=True,
        )

    print(f"\ndone in {(time.time()-t0)/60:.1f} min | {len(list(OUT_DIR.glob('*.npz')))}/371 clips extracted")
    return 0


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