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   # detection feature index where keypoints begin
_KPT_STRIDE = 3  # (x, y, conf) per keypoint


def _yolo_to_keypoints(out: np.ndarray, conf_thresh: float = 0.25, iou_thresh: float = 0.5):
    """out shape: (1, N+4, num_detections) — 84 for det-only, 51+5 for pose."""
    # (1, C, D) -> (D, C) so each row is a candidate detection.
    preds = out[0].T
    conf = preds[:, 4]
    keep = conf > conf_thresh
    preds = preds[keep]
    if not len(preds):
        return []

    # NMS via greedy diagonal covariance check (good enough at pose level).
    def _nms(rows, thresh):
        order = conf[keep].argsort()[::-1]
        keep_idx = []
        suppressed = set()
        centers = rows[:, :2]
        wh = rows[:, 2:4]
        areas = wh[:, 0] * wh[:, 1]
        for i in order:
            if i in suppressed:
                continue
            keep_idx.append(i)
            cx1, cy1 = centers[i]
            w1, h1 = wh[i]
            for j in order:
                if j == i or j in suppressed:
                    continue
                cx2, cy2 = centers[j]
                w2, h2 = wh[j]
                # Intersection over smaller area (proxy for pose NMS).
                dx = max(0, min(cx1 + w1/2, cx2 + w2/2) - max(cx1 - w1/2, cx2 - w2/2))
                dy = max(0, min(cy1 + h1/2, cy2 + h2/2) - max(cy1 - h1/2, cy2 - h2/2))
                inter = dx * dy
                smaller = min(areas[i], areas[j])
                if smaller > 0 and inter / smaller > thresh:
                    suppressed.add(j)
        return rows[keep_idx]

    preds = _nms(preds, iou_thresh)
    kpts = []
    for row in preds:
        kp = row[_KPT_START:].reshape(-1, _KPT_STRIDE)
        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.

    Uses shoulder-width normalisation so tags stay pose-invariant (scale blur).
    `kp[:, 0]` = x, `kp[:, 1]` = y, `kp[:, 2]` = confidence per keypoint.
    """
    tags: list[str] = []
    if not kpts:
        return tags

    for kp in kpts:  # kp: (17, 3), COCO order
        # Shoulder width — pose-invariant scale reference.
        ls, rs = kp[5], kp[6]
        sw = max(abs(ls[0] - rs[0]), 1e-4)
        lh, rh = kp[11], kp[12]
        th = max(abs(lh[0] - rh[0]), 1e-4)

        lw, rw = kp[9], kp[10]
        lk, rk = kp[13], kp[14]
        la, ra = kp[15], kp[16]

        # ---- arms ----
        for side, w, s in (("left", lw, ls), ("right", rw, rs)):
            if w[2] < 0.3 or s[2] < 0.3:
                continue
            # arm up: wrist significantly above shoulder
            if w[1] < s[1] - 0.15 * sw:
                tags.append(f"{side} arm up")
            # arm extended sideways: wrist is far from shoulder laterally
            elif abs(w[0] - s[0]) > 0.4 * sw:
                tags.append(f"{side} arm out")

        # ---- posture: torso vs limb geometry ----
        if lh[2] > 0.3 and lk[2] > 0.3:
            # sitting: knee is raised toward hip level
            if lk[1] < lh[1] - 0.05 * th:
                tags.append("sitting")
            else:
                tags.append("standing")

        # lying: torso is more horizontal than vertical
        if ls[2] > 0.3 and lh[2] > 0.3:
            tv = abs(ls[1] - lh[1])
            th_ = abs(ls[0] - lh[0])
            if th_ > tv * 1.5:
                tags.append("lying")

        # walking/running: legs scissor horizontally
        if la[2] > 0.3 and ra[2] > 0.3:
            stride = abs(la[0] - ra[0])
            if stride > 0.8 * sw:
                tags.append("walking" if stride < 2.0 * sw else "running")

        # kneeling: both knees below hips significantly, sitting
        if lk[2] > 0.3 and rk[2] > 0.3 and lh[2] > 0.3:
            if lk[1] > lh[1] + 0.4 * th and rk[1] > lh[1] + 0.4 * th:
                tags.append("kneeling")

        # torso lean: shoulder-to-hip vector strongly non-vertical
        if ls[2] > 0.3 and lh[2] > 0.3:
            lean = abs(ls[0] - lh[0]) / max(abs(ls[1] - lh[1]), 1e-4)
            if lean > 0.35:
                tags.append("leaning")

    # Deduplicate — for multi-person detections keep only the majority vote.
    from collections import Counter
    counts = Counter(tags)
    return [t for t, n in sorted(counts.items(), key=lambda p: (-p[1], p[0]))]



# ---------------------------------------------------------------------------
# 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 = "Xenova/yolov8n-pose", filename: str = "onnx/model.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