src/pose_tagger.py

"""Pose estimation from an image via DWPose/YOLO11n ONNX Runtime.

Runs a lightweight YOLO11n-pose ONNX model for person detection + 17-keypoint
pose estimation. The 17 COCO keypoints are turned into a natural-language
pose tag string that can be appended to a prompt (e.g. ``standing, arms up``).
"""
from __future__ import annotations

import json
import os
from typing import Optional

import numpy as np
from PIL import Image, ImageOps

try:
    import onnxruntime as ort
except Exception:  # pragma: no cover
    ort = None

# ---------------------------------------------------------------------------
# Keypoint analysis
# ---------------------------------------------------------------------------

import os

_DEBUG_POSE = os.environ.get("WHYX_DEBUG_POSE", "0").strip().lower() in ("1", "true", "yes")
_COCO_KP = [
    "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",
]

# YOLOv8-pose ONNX model keypoint order is COCO 17. Each detection row is
# [x, y, w, h, conf, kp0_x, kp0_y, kp0_conf, ..., kp16_x, kp16_y, kp16_conf].
_KPT_START = 5    # index of first keypoint value in a row of the output vector
_KPT_STRIDE = 3   # (x, y, conf) per keypoint


def _yolo_to_keypoints(out: np.ndarray, conf_thresh: float = 0.25):
    """out shape: (1, N+5, num_detections). Return list of (K, 3) arrays."""
    # (1, N+5, D) -> (D, N+5)
    preds = out[0].T
    conf = preds[:, 4]
    keep = conf > conf_thresh
    preds = preds[keep]
    kpts = []
    for row in preds:
        kp = row[_KPT_START:].reshape(-1, _KPT_STRIDE)  # (K, 3)
        kpts.append(kp)
    return kpts


def _keypoints_to_pose_tags(kpts: list[np.ndarray]) -> list[str]:
    """Heuristic mapping of COCO-17 coordinates to coarse pose tags."""
    tags: list[str] = []
    if not kpts:
        return tags

    for kp in kpts:
        # Normalise coords by shoulder width so the tags are pose-invariant.
        ls, rs = kp[5], kp[6]
        shoulder_span = max(np.abs(ls[0] - rs[0]), 1e-4)
        lw, rw = kp[9], kp[10]
        hips = kp[11], kp[12]
        lk, rk = kp[13], kp[14]
        la, ra = kp[15], kp[16]

        # arms up (wrists above shoulders)
        if lw[2] > 0.3 and rw[2] > 0.3 and lw[1] < ls[1] - 0.05 * shoulder_span \
                and rw[1] < rs[1] - 0.05 * shoulder_span:
            tags.append("arms up")
        # one arm up
        elif (lw[2] > 0.3 and lw[1] < ls[1] - 0.05 * shoulder_span) or (
                rw[2] > 0.3 and rw[1] < rs[1] - 0.05 * shoulder_span):
            tags.append("arm up")

        # sitting vs standing: hip-to-knee vertical distance smaller than
        # hip-to-ankle distance suggests knees bent (sitting)
        if hips[2] > 0.3 and lk[2] > 0.3:
            hip_y = hips[1]
            knee_y = lk[1]
            ankle_y = la[1] if la[2] > 0.3 else knee_y
            # Sitting: knees are raised toward hips.
            if knee_y < hip_y - 0.05 * shoulder_span:
                tags.append("sitting")
            else:
                tags.append("standing")

        # lying: body's vertical extent is smaller than horizontal
        if hips[2] > 0.3 and ls[2] > 0.3:
            torso_vert = abs(ls[1] - hips[1])
            torso_horiz = abs(ls[0] - hips[0])
            if torso_horiz > torso_vert * 1.5:
                tags.append("lying")

        # walking: one ankle significantly ahead of the other horizontally
        if la[2] > 0.3 and ra[2] > 0.3:
            stride = abs(la[0] - ra[0])
            if stride > 0.4 * shoulder_span:
                tags.append("walking")

    # Deduplicate preserving order
    seen: set[str] = set()
    out: list[str] = []
    for t in tags:
        if t not in seen:
            seen.add(t)
            out.append(t)
    return out


# ---------------------------------------------------------------------------
# ONNX pose runner
# ---------------------------------------------------------------------------


def _providers() -> list[str]:
    avail = ort.get_available_providers()
    return ["CUDAExecutionProvider", "CPUExecutionProvider"] if "CUDAExecutionProvider" in avail else ["CPUExecutionProvider"]


class PoseEstimator:
    """YOLO11n-pose wrapper with lazy model download."""

    def __init__(self, repo_id: str = "SamTheDev/YOLO11n-pose", filename: str = "yolo11n.onnx"):
        self._repo = repo_id
        self._filename = filename
        self._session: Optional["ort.InferenceSession"] = None
        self._input_shape: tuple[int, int] = (640, 640)
        self._loaded = False

    def ensure_loaded(self) -> bool:
        if self._loaded:
            return True
        if ort is None:
            return False
        from huggingface_hub import hf_hub_download
        try:
            model_path = hf_hub_download(repo_id=self._repo, filename=self._filename)
        except Exception:
            return False
        sess = ort.InferenceSession(model_path, providers=_providers())
        inp = sess.get_inputs()[0]
        shape = inp.shape
        if len(shape) == 4:
            self._input_shape = (int(shape[2]), int(shape[3]))
        self._session = sess
        self._loaded = True
        return True

    def estimate(self, image) -> dict:
        """Return {pose_tags, pose_score, people_count, raw_keypoints_count}."""
        if not self.ensure_loaded():
            return {"pose_tags": [], "people_count": 0}

        pil = image if isinstance(image, Image.Image) else Image.fromarray(np.asarray(image))
        pil = ImageOps.exif_transpose(ImageOps.fit(pil, self._input_shape, Image.LANCZOS))
        arr = np.asarray(pil, dtype=np.float32) / 255.0
        arr = arr.transpose(2, 0, 1)[None, ...]  # NCHW

        inp_name = self._session.get_inputs()[0].name
        out = self._session.run(None, {inp_name: arr})[0]
        kpts = _yolo_to_keypoints(out)
        if _DEBUG_POSE:
            print(f"[pose] raw_out={out.shape}, n_kpts_above_thresh={len(kpts)}")

        pose_tags = _keypoints_to_pose_tags(kpts)
        confs = [kp[:, 2].mean() for kp in kpts if kp.size]
        pose_score = float(np.mean(confs)) if confs else 0.0

        return {
            "pose_tags": pose_tags,
            "people_count": len(kpts),
            "pose_score": round(pose_score, 4),
        }


_pose_instance: PoseEstimator | None = None


def get_pose_tagger() -> PoseEstimator:
    global _pose_instance
    if _pose_instance is None:
        _pose_instance = PoseEstimator()
    return _pose_instance