Spaces:
Running
Running
| """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 | |