"""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