formscout/config.py

"""
FormScout pipeline configuration.
All model IDs, thresholds, k-values, and feature flags live here.
No scattered literals elsewhere in the codebase.
"""
import os
from pathlib import Path

ROOT = Path(__file__).parent.parent

# ─── Model IDs ───────────────────────────────────────────────────────────────
_YOLO_DIR = ROOT / "checkpoints" / "yolo26"

POSE_MODELS: dict[str, dict] = {
    # ── MediaPipe (official Tasks API, local checkpoint) ───────────────────
    "MediaPipe-Pose — full (~9 MB, CPU-friendly)": {
        "backend": "mediapipe",
        "path": str(ROOT / "checkpoints" / "mediapipe" / "pose_landmarker_full.task"),
        "params_m": 4.2,
    },
    # ── YOLO26 (local checkpoints) ─────────────────────────────────────────
    "YOLO26n — nano (0.7M, fastest)": {
        "backend": "yolo",
        "path": str(_YOLO_DIR / "yolo26n-pose.pt"),
        "params_m": 0.7,
    },
    "YOLO26s — small (3.5M)": {
        "backend": "yolo",
        "path": str(_YOLO_DIR / "yolo26s-pose.pt"),
        "params_m": 3.5,
    },
    "YOLO26m — medium (9M)": {
        "backend": "yolo",
        "path": str(_YOLO_DIR / "yolo26m-pose.pt"),
        "params_m": 9.0,
    },
    "YOLO26l — large (25.9M)": {
        "backend": "yolo",
        "path": str(_YOLO_DIR / "yolo26l-pose.pt"),
        "params_m": 25.9,
    },
    "YOLO26x — extra-large (57.6M)": {
        "backend": "yolo",
        "path": str(_YOLO_DIR / "yolo26x-pose.pt"),
        "params_m": 57.6,
    },
    # ── Sapiens2 (Phase 3 — needs custom repo + detector, 308-kp Sociopticon) ─
    "Sapiens2-0.4B [Phase 3, ~1.6 GB]": {
        "backend": "sapiens2",
        "hf_id": "facebook/sapiens2-pose-0.4b",
        "params_m": 400,
    },
    "Sapiens2-0.8B [Phase 3, ~3.2 GB]": {
        "backend": "sapiens2",
        "hf_id": "facebook/sapiens2-pose-0.8b",
        "params_m": 800,
    },
    "Sapiens2-1B [Phase 3, ~6 GB]": {
        "backend": "sapiens2",
        "hf_id": "facebook/sapiens2-pose-1b",
        "params_m": 1000,
    },
    "Sapiens2-5B [Phase 3, ~20 GB, large GPU]": {
        "backend": "sapiens2",
        "hf_id": "facebook/sapiens2-pose-5b",
        "params_m": 5000,
    },
}

DEFAULT_POSE_MODEL = "YOLO26n — nano (0.7M, fastest)"


def _is_model_available(spec: dict) -> bool:
    """Return True if the model checkpoint is present and the backend is importable."""
    backend = spec["backend"]
    if backend in ("yolo", "mediapipe"):
        return Path(spec["path"]).exists()
    if backend == "sapiens2":
        try:
            import sapiens  # noqa: F401  — custom repo must be installed
            return True
        except ImportError:
            return False
    return False


def available_pose_models() -> dict[str, dict]:
    """Subset of POSE_MODELS whose checkpoints/backends are actually ready."""
    return {name: spec for name, spec in POSE_MODELS.items() if _is_model_available(spec)}


# Backward-compat aliases
YOLO_POSE_MODEL = str(_YOLO_DIR / "yolo26l-pose.pt")
YOLO_POSE_MODEL_HQ = str(_YOLO_DIR / "yolo26x-pose.pt")
SAM_CHECKPOINT = "sam2.1_hiera_base_plus.pt"
SAM_3D_CHECKPOINT = ROOT / "checkpoints" / "sam-3d-body-dinov3" / "model.ckpt"
SAM_3D_HF_REPO = "facebook/sam-3d-body-dinov3"
SAM_3D_MHR_PATH = ROOT / "checkpoints" / "sam-3d-body-dinov3" / "assets" / "mhr_model.pt"
# ─── Judge / Classifier VLM (Qwen3-VL-8B-Instruct via llama.cpp) ────────────
# Default: stock Qwen3-VL-8B-Instruct Q4_K_M. To swap in a fine-tuned GGUF,
# set FORMSCOUT_JUDGE_GGUF (and FORMSCOUT_JUDGE_MMPROJ if it has its own
# projector) — no code change needed.
_QWEN_DIR = ROOT / "checkpoints" / "qwen3-vl"
JUDGE_GGUF = Path(os.environ.get(
    "FORMSCOUT_JUDGE_GGUF", _QWEN_DIR / "Qwen3VL-8B-Instruct-Q4_K_M.gguf"
))
JUDGE_MMPROJ = Path(os.environ.get(
    "FORMSCOUT_JUDGE_MMPROJ", _QWEN_DIR / "mmproj-Qwen3VL-8B-Instruct-F16.gguf"
))
JUDGE_HF_REPO = "Qwen/Qwen3-VL-8B-Instruct-GGUF"

QWEN_VLM_GGUF = str(JUDGE_GGUF)  # backward-compat alias
QWEN_EMBED_GGUF = "Qwen3-VL-Embedding-8B-Q4_K_M.gguf"
STGCN_CHECKPOINT = ROOT / "checkpoints" / "stgcn_fms.pth"

# ─── Pipeline flags ──────────────────────────────────────────────────────────
ENABLE_3D = False           # SAM 3D Body — access granted Jun 2026, off until integrated
ENABLE_STGCN = False        # Phase 3
ENABLE_RAG = False          # Phase 3
ENABLE_JUDGE = True         # VLM judge/classifier — falls back to rubric when llama-server is down

# ─── Thresholds ──────────────────────────────────────────────────────────────
MIN_CONFIDENCE = 0.6
SCORE_DISAGREE_THRESH = 1   # flag if |stgcn - judge| >= this
RETRIEVAL_K = 3

# ─── Video / Ingest ─────────────────────────────────────────────────────────
TARGET_FPS = 30.0
MAX_FRAMES = 300            # hard cap to avoid OOM
MAX_DURATION_SEC = 60.0     # warn on longer videos

# ─── Pose ────────────────────────────────────────────────────────────────────
POSE_BACKEND = "yolo"       # "yolo" | "sapiens"
POSE_CONF_THRESHOLD = 0.5
NUM_KEYPOINTS = 17

# ─── Biomechanics thresholds ────────────────────────────────────────────────
DEEP_SQUAT_FEMUR_HORIZONTAL_DEG = 90.0
DEEP_SQUAT_TORSO_TIBIA_MAX_DEG = 15.0
DEEP_SQUAT_KNEE_TRACKING_MARGIN_PX = 20

# ─── Serving (llama.cpp) ────────────────────────────────────────────────────
LLAMA_CPP_HOST = "127.0.0.1"
LLAMA_CPP_PORT_VLM = 8080
LLAMA_CPP_PORT_EMBED = 8081

# ─── Judge backend selection ────────────────────────────────────────────────
# "llama_cpp"   — local llama-server (default for local dev; works perfectly)
# "transformers"— in-process Qwen3-VL via transformers, GPU on HF Spaces (ZeroGPU)
# "auto"        — transformers ONLY on a GPU/ZeroGPU Space, else llama_cpp
JUDGE_BACKEND = os.environ.get("FORMSCOUT_JUDGE_BACKEND", "auto")
JUDGE_HF_MODEL = os.environ.get("FORMSCOUT_JUDGE_HF_MODEL", "Qwen/Qwen3-VL-8B-Instruct")
ON_HF_SPACE = bool(os.environ.get("SPACE_ID"))

# Seconds the ZeroGPU window stays allocated per analysis. One window wraps the
# whole pipeline (pose, optional 3D, Qwen3-VL judge), so size it for the slowest
# clip; raise via env for long videos. Only effective on a ZeroGPU Space.
ZEROGPU_DURATION = int(os.environ.get("FORMSCOUT_ZEROGPU_DURATION", "120"))


def has_gpu() -> bool:
    """True on a ZeroGPU Space (env flag) or when CUDA is actually present.

    ZeroGPU exposes no CUDA outside @spaces.GPU, so it is detected via the
    SPACES_ZERO_GPU env flag; ordinary GPU Spaces report via torch.cuda.
    """
    if os.environ.get("SPACES_ZERO_GPU") or os.environ.get("ZERO_GPU"):
        return True
    try:
        import torch
        return bool(torch.cuda.is_available())
    except Exception:
        return False


def resolve_judge_backend() -> str:
    """Resolve the effective judge backend from JUDGE_BACKEND + environment.

    `auto` only engages the heavy in-process transformers model when a GPU is
    actually available — a CPU-only Space stays on llama_cpp (which is then
    unreachable, so the Judge falls back to the fast rubric instead of trying to
    run a 17 GB model on CPU).
    """
    if JUDGE_BACKEND in ("llama_cpp", "transformers"):
        return JUDGE_BACKEND
    return "transformers" if (ON_HF_SPACE and has_gpu()) else "llama_cpp"