experiment-2107/scripts/config.py

"""Per-dataset pipeline configuration.

BASE_CONFIG holds settings shared by every source (this is the same dict pipeline.py used to
define inline). SOURCE_OVERRIDES layers per-dataset differences on top of it -- e.g. fs_jump3d's
mocap-derived clips shouldn't be frame-capped the way skatingverse's video extraction is.
get_source_config() merges the two for a given source name, using whatever config dict the
caller is actually running with (not always BASE_CONFIG) as the base, so runtime overrides set
by callers like run_500.py still take effect.
"""

from __future__ import annotations

from pathlib import Path

PROJECT_ROOT = Path(__file__).resolve().parents[2]

DATASET_PATHS = {
    "skatingverse": PROJECT_ROOT / "data" / "skatingverse",
    "current": PROJECT_ROOT / "data" / "current",
    "mmfs": PROJECT_ROOT / "data" / "mmfs",
    "fs_jump3d": PROJECT_ROOT / "data" / "fs_jump3d",
}

BASE_CONFIG = {
    "datasets": dict(DATASET_PATHS),
    "pose_estimator": "yolo",  # "mediapipe" (CPU) or "yolo" (GPU, see yolo_* keys below)
    "yolo_weights": "yolo11n-pose.pt",  # or "yolo11s-pose.pt" / "yolo11m-pose.pt"
    "yolo_imgsz": 256,  # 640 needed for reliable skater detection on wide rink shots (256 underdetects)
    "yolo_half": True,  # fp16 inference
    "yolo_max_batch": 128,  # cap frames per predict() call to bound VRAM under parallelism
    "yolo_decode_workers": 16,  # GIL-releasing cv2 decode threads feeding the single GPU consumer
    "target_sequence_length": 128,
    "confidence_threshold": 0.3,
    "smoothing_window": 7,
    "smoothing_polyorder": 3,
    "fps": 30.0,
    "max_videos": 20000,
    "default_width": 16.0,
    "default_height": 9.0,
    "train_ratio": 0.8,
    "val_ratio": 0.1,
    "test_ratio": 0.1,
    "random_seed": 42,
    "output_dir": PROJECT_ROOT / "data" / "processed_new",
    "skeleton_cache_dir": PROJECT_ROOT / "data" / "processed_new" / "skeleton_cache",
    "num_workers": 12,
    "video_num_workers": 1,
    "suppress_mediapipe_logs": True,
    "mediapipe_model_complexity": 0,
    "extract_every_n_frames": 1,
    "max_video_frames": 300,
    "show_frame_progress": False,
}

# Per-source overrides layered on top of the running config. Only list keys that differ from
# whatever the caller is already using -- everything not mentioned here is inherited untouched.
SOURCE_OVERRIDES: dict[str, dict] = {
    "fs_jump3d": {
        # FS-Jump3D's samples come from its pre-extracted 3D mocap (npy/), not from running
        # pose estimation over the raw video -- so no frame budget is being spent on estimation
        # for this source today, and there's no reason to cap it. Left uncapped "for now"; revisit
        # if fs_jump3d's own 12-camera RGB video/ ever gets wired into the video-extraction path
        # too, since that path (extract_skeleton_from_video / gpu_pose) does read this key.
        "max_video_frames": None,
        # Confirmed from the raw per-clip JSON's Timebase.Frequency (60.0 across every sample
        # checked) -- the mocap system captures at 60Hz, not the 30fps default tuned for the
        # other (video) sources. Getting this right matters because process_sample derives
        # velocity/angular-velocity features from fps.
        "fps": 60.0,
    },
}


def get_source_config(source: str, base: dict) -> dict:
    """Return a config dict for one source: `base` (the config actually being run) with that
    source's overrides layered on top. `datasets` is shallow-copied so per-source mutation
    can't leak between sources sharing one run_pipeline() call."""
    merged = dict(base)
    merged["datasets"] = dict(base["datasets"])
    merged.update(SOURCE_OVERRIDES.get(source, {}))
    return merged