formscout/agents/pose2d.py

"""
Pose2DAgent — 2D per-frame keypoint extraction.

Backends: yolo (local checkpoints, ultralytics), mediapipe (official Tasks API,
          local .task checkpoint), sapiens2 (Meta HF/transformers).
All backends output COCO-17 keypoints: dict[int, {x, y, conf}] per frame.

Input:  IngestResult
Output: Pose2DResult(keypoints per frame, fps, confidence)
Failure: Pose2DResult(confidence=0.0, notes=<reason>) — never raises.
Gated: yolo=no; mediapipe=no (local checkpoint); sapiens2=yes (access accepted).
"""
from __future__ import annotations

import logging
import numpy as np

from formscout import config
from formscout.types import IngestResult, Pose2DResult

logger = logging.getLogger(__name__)

COCO_KEYPOINTS = [
    "nose", "left_eye", "right_eye", "left_ear", "right_ear",
    "left_shoulder", "right_shoulder", "left_elbow", "right_elbow",
    "left_wrist", "right_wrist", "left_hip", "right_hip",
    "left_knee", "right_knee", "left_ankle", "right_ankle",
]

# BlazePose-33 source indices → COCO-17 target indices
# BlazePose: 0=nose, 2=left_eye, 5=right_eye, 7=left_ear, 8=right_ear,
#            11=left_shoulder, 12=right_shoulder, 13=left_elbow, 14=right_elbow,
#            15=left_wrist, 16=right_wrist, 23=left_hip, 24=right_hip,
#            25=left_knee, 26=right_knee, 27=left_ankle, 28=right_ankle
_BP_SRC = [0, 2, 5, 7, 8, 11, 12, 13, 14, 15, 16, 23, 24, 25, 26, 27, 28]
_BP_DST = list(range(17))  # COCO indices 0..16

_model_cache: dict[str, object] = {}


# ── YOLO backend ──────────────────────────────────────────────────────────────

def _get_yolo(path: str) -> object:
    if path not in _model_cache:
        from ultralytics import YOLO
        _model_cache[path] = YOLO(path)
    return _model_cache[path]


def _run_yolo(frames: list, path: str) -> list[dict]:
    model = _get_yolo(path)
    out = []
    for frame in frames:
        try:
            results = model(frame, verbose=False)
            kps: dict[int, dict] = {}
            if results and results[0].keypoints is not None:
                kp = results[0].keypoints
                if kp.xy is not None and len(kp.xy) > 0:
                    xy = kp.xy[0].cpu().numpy()
                    conf = kp.conf[0].cpu().numpy()
                    for j in range(min(len(xy), 17)):
                        kps[j] = {"x": float(xy[j, 0]), "y": float(xy[j, 1]), "conf": float(conf[j])}
            out.append(kps)
        except Exception:
            out.append({})
    return out


# ── MediaPipe backend (official Tasks API, local .task checkpoint) ────────────

def _get_mediapipe_landmarker(path: str) -> object:
    """Return PoseLandmarker cached by model path."""
    cache_key = f"mp:{path}"
    if cache_key not in _model_cache:
        from mediapipe.tasks import python as mp_tasks
        from mediapipe.tasks.python import vision

        options = vision.PoseLandmarkerOptions(
            base_options=mp_tasks.BaseOptions(model_asset_path=path),
            running_mode=vision.RunningMode.IMAGE,
            num_poses=1,
            min_pose_detection_confidence=0.4,
            min_pose_presence_confidence=0.4,
            min_tracking_confidence=0.4,
        )
        _model_cache[cache_key] = vision.PoseLandmarker.create_from_options(options)
    return _model_cache[cache_key]


def _run_mediapipe(frames: list, path: str) -> list[dict]:
    import cv2
    import mediapipe as mp

    try:
        landmarker = _get_mediapipe_landmarker(path)
    except Exception as e:
        logger.warning("mediapipe load failed: %s", e)
        return [{} for _ in frames]

    out = []
    for frame in frames:
        try:
            h, w = frame.shape[:2]
            rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            mp_image = mp.Image(image_format=mp.ImageFormat.SRGB, data=rgb)
            detection = landmarker.detect(mp_image)

            kps: dict[int, dict] = {}
            if detection.pose_landmarks:
                lms = detection.pose_landmarks[0]
                for coco_idx, bp_idx in zip(_BP_DST, _BP_SRC):
                    if bp_idx < len(lms):
                        lm = lms[bp_idx]
                        kps[coco_idx] = {
                            "x": float(lm.x * w),
                            "y": float(lm.y * h),
                            "conf": float(lm.visibility),
                        }
            out.append(kps)
        except Exception:
            out.append({})
    return out


# ── Sapiens2 backend (Meta HF, transformers) ──────────────────────────────────

def _get_sapiens2(hf_id: str) -> object:
    if hf_id not in _model_cache:
        from transformers import pipeline as hf_pipeline
        _model_cache[hf_id] = hf_pipeline("pose-estimation", model=hf_id)
    return _model_cache[hf_id]


def _run_sapiens2(frames: list, hf_id: str) -> list[dict]:
    try:
        pipe = _get_sapiens2(hf_id)
    except Exception as e:
        logger.warning("sapiens2 load failed: %s", e)
        return [{} for _ in frames]

    from PIL import Image

    out = []
    for frame in frames:
        try:
            pil_img = Image.fromarray(frame)
            result = pipe(pil_img)

            if not result:
                out.append({})
                continue

            # Take highest-confidence person (first result)
            person = result[0]
            keypoints = person.get("keypoints", [])
            scores = person.get("keypoint_scores", [])

            # Build name→(x, y, score) lookup from pipeline output
            kp_lookup: dict[str, tuple] = {}
            for i, kp in enumerate(keypoints):
                if isinstance(kp, dict):
                    name = kp.get("label", "")
                    x, y = kp.get("x", 0.0), kp.get("y", 0.0)
                else:
                    name = ""
                    x, y = float(kp[0]), float(kp[1])
                score = float(scores[i]) if i < len(scores) else 0.0
                if name:
                    kp_lookup[name] = (x, y, score)

            kps: dict[int, dict] = {}
            for coco_idx, name in enumerate(COCO_KEYPOINTS):
                if name in kp_lookup:
                    x, y, s = kp_lookup[name]
                    kps[coco_idx] = {"x": x, "y": y, "conf": s}
            out.append(kps)
        except Exception:
            out.append({})
    return out


# ── Agent ─────────────────────────────────────────────────────────────────────

class Pose2DAgent:
    """Extracts COCO-17 keypoints per frame; dispatches to YOLO, MediaPipe, or Sapiens2."""

    def run(self, ingest: IngestResult, model_key: str | None = None) -> Pose2DResult:
        if not ingest.frames:
            return Pose2DResult(keypoints=[], fps=ingest.fps, confidence=0.0, notes="no frames in ingest")

        key = model_key or config.DEFAULT_POSE_MODEL
        spec = config.POSE_MODELS.get(key)
        if spec is None:
            logger.warning("Unknown model_key %r — falling back to %s", key, config.DEFAULT_POSE_MODEL)
            spec = config.POSE_MODELS[config.DEFAULT_POSE_MODEL]

        backend = spec["backend"]
        try:
            if backend == "yolo":
                kps_per_frame = _run_yolo(ingest.frames, spec["path"])
            elif backend == "mediapipe":
                kps_per_frame = _run_mediapipe(ingest.frames, spec["path"])
            elif backend == "sapiens2":
                kps_per_frame = _run_sapiens2(ingest.frames, spec["hf_id"])
            else:
                return Pose2DResult(
                    keypoints=[{} for _ in ingest.frames],
                    fps=ingest.fps, confidence=0.0,
                    notes=f"unknown backend: {backend}",
                )
        except Exception as e:
            return Pose2DResult(
                keypoints=[{} for _ in ingest.frames],
                fps=ingest.fps, confidence=0.0,
                notes=str(e),
            )

        n_detected = sum(1 for f in kps_per_frame if f)
        total_conf = sum(
            sum(kp["conf"] for kp in f.values()) / len(f)
            for f in kps_per_frame if f
        )
        overall_conf = (total_conf / n_detected) if n_detected > 0 else 0.0
        notes = "" if n_detected > 0 else "no person detected in any frame"

        return Pose2DResult(
            keypoints=kps_per_frame,
            fps=ingest.fps,
            confidence=overall_conf,
            notes=notes,
        )