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