experiment-2107/scripts/temporal_scripts/fs_spot_data.py

"""FS-Spot: untrimmed figure skating data with frame-precise jump/spin boundaries.

Source: Hong et al., "Spotting Temporally Precise, Fine-Grained Events in Video"
(ECCV 2022) -- https://github.com/jhong93/spot. The repo ships ONLY annotations
(train/val/test.json per split, committed to git); the broadcast video itself is
ISU/Olympic copyright and is not redistributed, so videos.csv gives YouTube ids
and you source the pixels yourself. Everything is indexed in frames at 25fps, so
the download MUST be the 25fps rendition (format 248) or every label shifts.

This module turns that into skeleton clips our pipeline can train on:

    11 broadcasts (4-6h each, 29-40 skaters back to back)
      -> 371 per-skater clips (frame windows carved out by the clip name)
      -> YOLO pose per frame, skater picked by a continuity tracker
      -> (T, 17, 3) COCO-17 normalized skeletons + dense per-frame labels

Why this dataset: it is the only public figure skating source with untrimmed
multi-action sequences AND ground-truth temporal boundaries. 1,464 jumps, each
with a takeoff frame, a landing frame, and a jump type; 373 spins; 338 real jump
combinations. The 6 jump types are exactly FS-Jump3D's taxonomy.

Clip naming encodes the window: <broadcast>_<skater##>_<startframe>_<endframe>,
where endframe - startframe == num_frames (verified across all 371).
"""

from __future__ import annotations

import json
import re
import subprocess
from dataclasses import dataclass, field
from pathlib import Path

import numpy as np

PROJECT_ROOT = Path(__file__).resolve().parents[2]
SPOT_ANNO_DIR = PROJECT_ROOT / "data" / "fs_spot" / "spot_annotations"
VIDEO_DIR = PROJECT_ROOT / "data" / "fs_spot" / "videos"

# Most broadcasts download as VP9 (format 248, .webm). Two have no VP9 rendition
# and fall back to format 137 (h264, .mp4) -- same 1920x1080 @ 25fps, so frame
# indices are unaffected. Look up by stem, not by a hardcoded extension.
VIDEO_EXTS = (".webm", ".mp4", ".mkv")
OUT_DIR = PROJECT_ROOT / "data" / "fs_spot" / "processed"

FPS = 25.0
CLIP_RE = re.compile(r"^(?P<broadcast>.+)_(?P<skater>\d+)_(?P<start>\d{8})_(?P<end>\d{8})$")

# spot's `comment` field carries the jump type. One record is typo'd 'tow_loop'.
JUMP_TYPE_FIX = {"tow_loop": "toe_loop"}

# Dense per-frame taxonomy. Jump types match FS_JUMP3D_TAXONOMY names so a model
# trained here transfers to the FS-Jump3D mocap clips without a label remap.
SPOT_TAXONOMY = {
    0: "NONE",
    1: "Axel",
    2: "Flip",
    3: "Loop",
    4: "Lutz",
    5: "Salchow",
    6: "Toeloop",
    7: "Spin",
}
SPOT_LABEL_TO_IDX = {name: idx for idx, name in SPOT_TAXONOMY.items()}

_SPOT_JUMP_TO_TAXONOMY = {
    "axel": "Axel",
    "flip": "Flip",
    "loop": "Loop",
    "lutz": "Lutz",
    "salchow": "Salchow",
    "toe_loop": "Toeloop",
}


@dataclass
class Element:
    """One jump or spin: a [takeoff, landing] frame window, relative to the clip."""

    kind: str  # "jump" | "spin"
    label: str  # SPOT_TAXONOMY name
    takeoff: int
    landing: int

    @property
    def duration(self) -> int:
        return self.landing - self.takeoff


@dataclass
class Clip:
    """One skater's performance: a frame window inside a broadcast."""

    name: str
    broadcast: str
    start: int  # absolute frame in the broadcast
    end: int
    num_frames: int
    elements: list[Element] = field(default_factory=list)

    @property
    def transcript(self) -> list[str]:
        """Ordered element labels -- the CTC target, if you want one."""
        return [e.label for e in self.elements]

    def frame_labels(self) -> np.ndarray:
        """Dense (T,) int labels: every frame in [takeoff, landing] gets the element
        label, everything else is NONE. This is the supervision CTC exists to avoid
        needing -- since we have it, prefer a dense loss and skip CTC entirely."""
        y = np.zeros(self.num_frames, dtype=np.int64)
        for e in self.elements:
            y[e.takeoff : e.landing + 1] = SPOT_LABEL_TO_IDX[e.label]
        return y


def _pair_events(events: list[dict]) -> list[Element]:
    """Pair takeoff/landing events into elements.

    spot stores events as a flat frame-ordered list of {frame, label, comment}.
    A takeoff is always followed by its landing, so pair them in order per kind.
    An unmatched takeoff at the end of a clip is dropped rather than guessed.

    Event frames are already CLIP-relative (0..num_frames), not absolute into the
    broadcast -- the clip name's start/end give the absolute window separately.
    """
    out: list[Element] = []
    pending: dict[str, tuple[int, str | None]] = {}
    for e in sorted(events, key=lambda x: x["frame"]):
        label = e["label"]
        kind, _, phase = label.partition("_")
        if phase == "takeoff":
            pending[kind] = (e["frame"], e.get("comment"))
        elif phase == "landing":
            if kind not in pending:
                continue  # landing with no takeoff -- skip rather than invent one
            takeoff, comment = pending.pop(kind)
            if kind == "jump":
                raw = JUMP_TYPE_FIX.get(comment or "", comment or "")
                name = _SPOT_JUMP_TO_TAXONOMY.get(raw)
                if name is None:
                    continue  # unknown jump type -- don't silently mislabel it
            else:
                name = "Spin"
            out.append(Element(kind=kind, label=name, takeoff=takeoff, landing=e["frame"]))
    return out


def build_index(anno_dir: Path = SPOT_ANNO_DIR, split_set: str = "fs_comp") -> dict[str, Clip]:
    """Parse spot's json into Clip records keyed by clip name.

    fs_comp and fs_perf are two different train/val/test *partitions of the same
    371 clips*, not different data -- so indexing either gives the full set. We
    use fs_comp's partition by default.
    """
    clips: dict[str, Clip] = {}
    for split in ("train", "val", "test"):
        path = anno_dir / split_set / f"{split}.json"
        for rec in json.loads(path.read_text()):
            m = CLIP_RE.match(rec["video"])
            if m is None:
                raise ValueError(f"unparseable clip name: {rec['video']!r}")
            start, end = int(m["start"]), int(m["end"])
            if end - start != rec["num_frames"]:
                raise ValueError(
                    f"{rec['video']}: name span {end - start} != num_frames {rec['num_frames']}"
                )
            clips[rec["video"]] = Clip(
                name=rec["video"],
                broadcast=m["broadcast"],
                start=start,
                end=end,
                num_frames=rec["num_frames"],
                elements=_pair_events(rec["events"]),
            )
    return clips


def find_video(broadcast: str, video_dir: Path = VIDEO_DIR) -> Path | None:
    """Locate a broadcast's file regardless of container."""
    for ext in VIDEO_EXTS:
        p = video_dir / f"{broadcast}{ext}"
        if p.exists():
            return p
    return None


def available_broadcasts(video_dir: Path = VIDEO_DIR) -> set[str]:
    """Broadcast stems that have a complete (non-.part) file on disk."""
    return {p.stem for p in video_dir.iterdir() if p.suffix in VIDEO_EXTS}


def _ffmpeg_window_cmd(
    video: Path, start_frame: int, num_frames: int, scale: tuple[int, int] | None = None
) -> list[str]:
    """Decode exactly [start_frame, start_frame+num_frames) to rawvideo on stdout.

    Input-side -ss makes ffmpeg seek to the preceding keyframe and then decode
    forward to the exact frame (accurate seek). That precision is the whole ball
    game here: every label is a frame index, so a keyframe-snapped cut would shift
    every annotation in the clip. Validated against a known Axel -- the skater
    leaves the ice within 2 frames of the annotated takeoff.

    `scale` downsizes in ffmpeg. This is free accuracy-wise: YOLO letterboxes to
    its imgsz regardless, so 1920x1080 and 640x360 both arrive at the model as
    640x360 -- but pre-scaling cuts ~9x off the decode/transfer bytes and moves
    the letterbox off ultralytics' single-threaded CPU path (78 -> ~300 fps).
    Keypoints are normalized to [0,1] by frame size, so coordinates are unchanged.
    """
    cmd = [
        "ffmpeg", "-v", "error", "-nostdin",
        "-ss", f"{start_frame / FPS:.6f}",
        "-i", str(video),
        "-frames:v", str(num_frames),
    ]
    if scale is not None:
        cmd += ["-vf", f"scale={scale[0]}:{scale[1]}"]
    return cmd + ["-f", "rawvideo", "-pix_fmt", "rgb24", "-"]


def decode_window(video: Path, start_frame: int, num_frames: int, width: int, height: int) -> np.ndarray:
    """Decode a window fully into memory as (N, H, W, 3) uint8.

    Only safe for short windows: 1080p costs ~6.2MB/frame, so a full 4,266-frame
    clip would be ~26GB. Use stream_window() for anything clip-length.
    """
    proc = subprocess.run(_ffmpeg_window_cmd(video, start_frame, num_frames), capture_output=True)
    if proc.returncode != 0:
        raise RuntimeError(f"ffmpeg failed on {video.name}@{start_frame}: {proc.stderr.decode()[:400]}")
    frame_bytes = width * height * 3
    n = len(proc.stdout) // frame_bytes
    if n == 0:
        raise RuntimeError(f"ffmpeg returned no frames for {video.name}@{start_frame}")
    return np.frombuffer(proc.stdout[: n * frame_bytes], dtype=np.uint8).reshape(n, height, width, 3)


def stream_window(
    video: Path,
    start_frame: int,
    num_frames: int,
    width: int,
    height: int,
    chunk: int = 128,
    scale: tuple[int, int] | None = None,
):
    """Yield (N, H, W, 3) uint8 chunks over the window, never holding it all at once.

    A clip is ~4,266 frames of 1080p (~26GB decoded), so the extractor consumes it
    in chunks and keeps only the (T, 17, 3) skeletons (~1.7MB) -- four orders of
    magnitude smaller than the pixels they came from.

    `width`/`height` must match what ffmpeg emits, i.e. `scale` when it is set.
    """
    frame_bytes = width * height * 3
    proc = subprocess.Popen(
        _ffmpeg_window_cmd(video, start_frame, num_frames, scale=scale),
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        bufsize=frame_bytes * 4,
    )
    try:
        remaining = num_frames
        while remaining > 0:
            want = min(chunk, remaining)
            buf = proc.stdout.read(frame_bytes * want)
            if not buf:
                break
            got = len(buf) // frame_bytes
            if got == 0:
                break
            yield np.frombuffer(buf[: got * frame_bytes], dtype=np.uint8).reshape(got, height, width, 3)
            remaining -= got
    finally:
        if proc.stdout:
            proc.stdout.close()
        proc.wait()
        if proc.returncode not in (0, None):
            err = proc.stderr.read().decode()[:400] if proc.stderr else ""
            raise RuntimeError(f"ffmpeg failed on {video.name}@{start_frame}: {err}")
        if proc.stderr:
            proc.stderr.close()


def pick_skater(boxes_xyxy: np.ndarray, confs: np.ndarray, prev_center: np.ndarray | None) -> int:
    """Choose which detected person is the skater.

    gpu_pose._results_to_skeleton takes argmax(conf), which is fine for trimmed
    single-skater clips but wrong here: broadcast footage is full of crowd,
    coaches and judges, and a sharp face in the stands routinely outscores a
    motion-blurred skater. Score on confidence AND size (the skater is the camera
    subject, so they're the large one), then bias toward the previous frame's pick
    so the choice stays on one person through a crowd pan.
    """
    areas = (boxes_xyxy[:, 2] - boxes_xyxy[:, 0]) * (boxes_xyxy[:, 3] - boxes_xyxy[:, 1])
    score = confs * np.sqrt(np.maximum(areas, 1.0))
    if prev_center is not None:
        centers = np.stack(
            [(boxes_xyxy[:, 0] + boxes_xyxy[:, 2]) / 2, (boxes_xyxy[:, 1] + boxes_xyxy[:, 3]) / 2], axis=1
        )
        dist = np.linalg.norm(centers - prev_center[None, :], axis=1)
        # Halve the score for every ~25% of frame width the candidate has jumped.
        score = score / (1.0 + (dist / 250.0) ** 2)
    return int(np.argmax(score))