"""Pose estimation from an image via YOLO-pose ONNX Runtime. Runs a YOLO-pose ONNX model for person detection + 17-keypoint pose estimation. The 17 COCO keypoints are mapped to Danbooru-style pose tags (e.g. ``standing, arms_up``) that blend into the WD14 general-tag list. Default checkpoint is ``Xenova/yolov8m-pose`` (~83 MB) — letterboxed at 1024x1024 for reliable anime character detection. Override via ``WHYX_POSE_MODEL`` / ``WHYX_POSE_FILENAME`` / ``WHYX_POSE_INPUT``. """ from __future__ import annotations import math 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 # --------------------------------------------------------------------------- _DEBUG_POSE = os.environ.get("WHYX_DEBUG_POSE", "0").strip().lower() in ("1", "true", "yes") _POSE_REPO = os.environ.get("WHYX_POSE_MODEL", "Xenova/yolov8m-pose") _POSE_FILENAME = os.environ.get("WHYX_POSE_FILENAME", "onnx/model.onnx") _POSE_INPUT = int(os.environ.get("WHYX_POSE_INPUT", "1024")) _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 # Wrist keypoint confidence gate for arm pose heuristics. Wrists are the most # jittery COCO keypoints; 0.30 produced false "one_arm_up" hits, 0.35 is the # sweet spot between recall and noise. _WRIST_MIN = 0.35 # Minimum mean keypoint confidence for a detection to emit pose tags. # Below this threshold the detection is considered unreliable and pose_tags # are suppressed (keypoints are still kept for skeleton overlay). _MIN_POSE_CONF = 0.35 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 _vis(kp: np.ndarray, i: int, thresh: float = 0.30) -> bool: """Keypoint visibility: confidence gate for COCO-17 row ``i``.""" return kp[i, 2] >= thresh def _mid(a, b): """Midpoint of two keypoint (x, y) pairs; None if both missing.""" if a is None and b is None: return None if a is None: return b if b is None: return a return ((a[0] + b[0]) / 2.0, (a[1] + b[1]) / 2.0) def _hypot(dx: float, dy: float) -> float: return float((dx * dx + dy * dy) ** 0.5) def _single_person_tags(kp: np.ndarray) -> list[str]: """Map one person's COCO-17 keypoints to Danbooru-style pose tags. All distances are normalised by shoulder width so the tags stay pose-invariant across image scales. Returns booru tags with underscores (e.g. ``arms_up``) so they blend into the WD14 general-tag list. Tag names are chosen to match the WD14 v3 vocabulary where one exists (``arms_at_sides``, ``praying``, ``crossed_legs``, ...); a few booru-only tags (``crouching``, ``on all fours``, ...) are emitted when no valid WD14 synonym exists. """ out: set[str] = set() ls, rs = kp[5], kp[6] le_, re_ = kp[7], kp[8] # elbows lw, rw = kp[9], kp[10] # wrists lh, rh = kp[11], kp[12] lk, rk = kp[13], kp[14] la, ra = kp[15], kp[16] nose = kp[0] lea, rea = kp[3], kp[4] # ears sw = max(abs(ls[0] - rs[0]), 1e-4) torso_ok = _vis(kp, 5) and _vis(kp, 6) and _vis(kp, 11) and _vis(kp, 12) sc = _mid(ls[:2], rs[:2]) hc = _mid(lh[:2], rh[:2]) kc = _mid(lk[:2], rk[:2]) ac = _mid(la[:2], ra[:2]) # ------------------------------------------------------------------ # Torso posture # ------------------------------------------------------------------ # On all fours: torso roughly horizontal, hips ABOVE the knees and both # wrists planted near the knees (hands and knees on the ground). all_fours = False if torso_ok and sc is not None and _vis(kp, 13) and _vis(kp, 14): if ( abs(lw[1] - lk[1]) < 0.4 * sw and abs(lw[0] - lk[0]) < 0.8 * sw and abs(rw[1] - rk[1]) < 0.4 * sw and abs(rw[0] - rk[0]) < 0.8 * sw ): all_fours = True if torso_ok and sc is not None and hc is not None: tvx, tvy = sc[0] - hc[0], sc[1] - hc[1] torso_len = _hypot(tvx, tvy) theta = abs(tvx) / max(torso_len, 1e-4) # 0=vertical, 1=horizontal if all_fours and theta > 0.55: out.add("on_all_fours") elif theta > 0.72: out.add("lying") elif kc is not None: knee_dy = kc[1] - hc[1] # + = knees below hips knee_dx = abs(kc[0] - hc[0]) # splits: both legs straight, spread very wide, hips near ground splits = ( _vis(kp, 13) and _vis(kp, 14) and _vis(kp, 15) and _vis(kp, 16) and abs(la[1] - lk[1]) < 0.25 * sw and abs(ra[1] - rk[1]) < 0.25 * sw and abs(la[0] - ra[0]) > 1.7 * sw and hc[1] > kc[1] - 0.2 * sw ) if splits: out.add("splits") elif knee_dy < -0.08 * sw: # knees raised toward the hips → seated out.add("sitting") elif ( _vis(kp, 15) and _vis(kp, 16) and la[1] <= lk[1] + 0.2 * sw and ra[1] <= rk[1] + 0.2 * sw and knee_dy < 0.35 * sw ): # both feet pulled up to knee level with knees raised → # airborne tuck jump out.add("jumping") elif knee_dy < 0.35 * sw and knee_dx < 0.8 * sw: # hips very close to knees with knees tucked out.add("squatting") if knee_dy < 0.15 * sw and knee_dx < 0.6 * sw: out.add("crouching") else: # kneeling: knees planted low, hips above, shins horizontal knee_planted_l = _vis(kp, 13) and _vis(kp, 15) and abs(lk[1] - la[1]) < 0.25 * sw and abs(lk[0] - la[0]) < 0.5 * sw knee_planted_r = _vis(kp, 14) and _vis(kp, 16) and abs(rk[1] - ra[1]) < 0.25 * sw and abs(rk[0] - ra[0]) < 0.5 * sw foot_planted_l = _vis(kp, 13) and _vis(kp, 15) and la[1] > lk[1] + 0.35 * sw foot_planted_r = _vis(kp, 14) and _vis(kp, 16) and ra[1] > rk[1] + 0.35 * sw one_knee = ( (knee_planted_l and foot_planted_r and not knee_planted_r) or (foot_planted_l and knee_planted_r and not knee_planted_l) ) kneeling = ( knee_planted_l and knee_planted_r and ac is not None and abs(ac[1] - kc[1]) < 0.25 * sw and abs(ac[0] - kc[0]) < 0.5 * sw ) if one_knee: out.add("kneeling_on_one_knee") elif kneeling: out.add("kneeling") else: # legs scissor → walking / running (mutually exclusive # with standing — no contradictory tags in the output). if _vis(kp, 15) and _vis(kp, 16) and _vis(kp, 13) and _vis(kp, 14): stride = abs(la[0] - ra[0]) if stride > 2.0 * sw: out.add("running") elif stride > 1.5 * sw: out.add("walking") else: out.add("standing") else: out.add("standing") else: out.add("standing") # torso lean / bent over (must not conflict with lying / all fours) if "lying" not in out and "on_all_fours" not in out and theta <= 0.72: if theta > 0.55: out.add("bent_over") elif theta > 0.3: out.add("leaning_forward" if tvx > 0 else "leaning_back") elif kc is not None and hc is not None: # partial visibility: infer sitting/standing from hips+knees only if kc[1] < hc[1] - 0.08 * sw: out.add("sitting") else: out.add("standing") # ------------------------------------------------------------------ # Arms (require both shoulders visible) # ------------------------------------------------------------------ if _vis(kp, 5) and _vis(kp, 6): sides = [] for w, e, s, h in ((lw, le_, ls, lh), (rw, re_, rs, rh)): if w[2] < _WRIST_MIN or s[2] < 0.3: continue side = {"w": w, "e": e, "s": s, "h": h} side["up"] = w[1] < s[1] - 0.15 * sw side["out"] = abs(w[0] - s[0]) > 0.4 * sw side["on_hip"] = abs(w[0] - h[0]) < 0.35 * sw and abs(w[1] - h[1]) < 0.3 * sw side["in_pocket"] = ( w[1] > h[1] + 0.1 * sw and w[1] < h[1] + 0.6 * sw and abs(w[0] - h[0]) < 0.3 * sw ) side["covers_face"] = ( _vis(kp, 0) and abs(w[0] - nose[0]) < 0.35 * sw and abs(w[1] - nose[1]) < 0.45 * sw ) side["hand_face"] = ( not side["covers_face"] and _vis(kp, 0) and abs(w[0] - nose[0]) < 0.5 * sw and nose[1] - 0.2 * sw < w[1] < nose[1] + 0.5 * sw ) # arm hanging straight down beside the body side["down"] = ( e[2] >= 0.3 and w[1] > e[1] + 0.05 * sw and w[1] > s[1] + 0.2 * sw and abs(w[0] - s[0]) < 0.5 * sw ) # wrist on the chest / belly (torso midline, guard sc/hc) side["on_chest"] = ( sc is not None and hc is not None and abs(w[0] - sc[0]) < 0.45 * sw and sc[1] + 0.25 * sw < w[1] < hc[1] + 0.05 * sw ) side["on_stomach"] = ( sc is not None and hc is not None and abs(w[0] - sc[0]) < 0.45 * sw and hc[1] - 0.15 * sw < w[1] < hc[1] + 0.45 * sw ) sides.append(side) for sd in sides: if sd["covers_face"]: out.add("covering_face") elif sd["hand_face"]: out.add("hand_on_own_face") hip_sides = [sd for sd in sides if sd["on_hip"] and not sd["up"]] pocket_sides = [sd for sd in sides if sd["in_pocket"] and not sd["up"]] n_up = sum(1 for sd in sides if sd["up"]) if n_up == 2: elbows_up = all(sd["e"][2] >= 0.3 and sd["e"][1] < sd["s"][1] + 0.05 * sw for sd in sides) near_head = _vis(kp, 0) and all(abs(sd["w"][0] - nose[0]) < 0.7 * sw for sd in sides) out.add("arms_behind_head" if (elbows_up and near_head) else "arms_up") elif n_up == 1: out.add("one_arm_up") # one arm straight overhead + the other hanging down → stretch up_side = next(sd for sd in sides if sd["up"]) others_down = all(sd["down"] for sd in sides if not sd["up"]) if ( others_down and up_side["w"][1] < up_side["s"][1] - 0.5 * sw and up_side["e"][2] >= 0.3 and up_side["w"][1] < up_side["e"][1] - 0.15 * sw ): out.add("stretching") # wrists close together: clasped hands → praying (chest) or arms # around neck (nape height) hands_together = ( len(sides) == 2 and abs(lw[0] - rw[0]) < 0.3 * sw and abs(lw[1] - rw[1]) < 0.3 * sw ) if hands_together and sc is not None and hc is not None: wmid_y = (lw[1] + rw[1]) / 2.0 if wmid_y < sc[1] + 0.35 * sw: out.add("arms_around_neck") elif wmid_y < hc[1] + 0.1 * sw: out.add("praying") # hands clasped at the small of the back arms_behind_back = ( len(sides) == 2 and sc is not None and hc is not None and abs(lw[0] - rw[0]) < 0.5 * sw and abs(lw[1] - rw[1]) < 0.35 * sw and lw[1] > hc[1] + 0.2 * sw and rw[1] > hc[1] + 0.2 * sw and abs(lw[0] - sc[0]) < 0.45 * sw and abs(rw[0] - sc[0]) < 0.45 * sw ) if arms_behind_back: out.add("arms_behind_back") # single hand on chest / belly (not when hands are clasped) if not hands_together: chest_sides = [ sd for sd in sides if sd["on_chest"] and not sd["covers_face"] and not sd["hand_face"] and not sd["on_hip"] and not sd["in_pocket"] and not sd["up"] ] if len(chest_sides) == 1: out.add("hand_on_own_chest") stomach_sides = [ sd for sd in sides if sd["on_stomach"] and not sd["covers_face"] and not sd["hand_face"] and not sd["on_hip"] and not sd["in_pocket"] and not sd["up"] ] if len(stomach_sides) == 1 and "arms_behind_back" not in out: out.add("hand_on_own_stomach") # both arms hanging straight down at the sides if ( len(sides) == 2 and all(sd["down"] for sd in sides) and not hands_together and not arms_behind_back and not any(sd["on_hip"] or sd["in_pocket"] for sd in sides) and "on_all_fours" not in out ): out.add("arms_at_sides") if len(sides) == 2: s0, s1 = sides # crossed arms: each wrist lands on the opposite side of the torso if ( s0["w"][0] > rs[0] and s1["w"][0] < ls[0] and all(sd["s"][1] < sd["w"][1] < sd["h"][1] for sd in sides) ): out.add("arms_crossed") elif len(hip_sides) == 2: out.add("hands_on_hips") elif len(pocket_sides) == 2: out.add("hands_in_pockets") elif all(sd["out"] and not sd["up"] for sd in sides): out.add("arms_out") if len(hip_sides) == 1: out.add("hand_on_hip") if len(pocket_sides) == 1: out.add("hand_in_pocket") # ------------------------------------------------------------------ # Head tilt (ears asymmetry) # ------------------------------------------------------------------ if _vis(kp, 3) and _vis(kp, 4) and abs(lea[1] - rea[1]) > 0.3 * sw: out.add("head_tilt") # ------------------------------------------------------------------ # Face orientation: frontal face (nose + both ears spread) → viewer # ------------------------------------------------------------------ if ( sc is not None and _vis(kp, 0) and _vis(kp, 3) and _vis(kp, 4) and abs(lea[0] - rea[0]) >= 0.35 * sw and nose[1] < sc[1] ): out.add("looking_at_viewer") # ------------------------------------------------------------------ # Legs detail (needs both ankles) # ------------------------------------------------------------------ if _vis(kp, 15) and _vis(kp, 16): stride = abs(la[0] - ra[0]) dy = abs(la[1] - ra[1]) # legs apart while standing if "standing" in out and stride > 1.2 * sw and "splits" not in out: out.add("legs_apart") # crossed legs while seated if "sitting" in out and stride < 0.3 * sw and dy < 0.25 * sw: out.add("crossed_legs") # ankles crossed while standing (one foot hooked past the other) elif ( "standing" in out and stride < 0.35 * sw and dy < 0.3 * sw and (la[0] > ra[0] + 0.1 * sw or ra[0] > la[0] + 0.1 * sw) ): out.add("crossed_legs") # one leg raised (foot clearly above the other, above its knee) if _vis(kp, 13) and _vis(kp, 14): hi, lo = (la, ra) if la[1] < ra[1] else (ra, la) hi_knee = lk if hi is la else rk if lo[1] - hi[1] > 0.5 * sw and hi[1] < hi_knee[1] - 0.1 * sw: out.add("one_leg_up") if "standing" in out: out.add("standing_on_one_leg") return list(out) def _keypoints_to_pose_tags(kpts: list[np.ndarray]) -> list[str]: """Heuristic mapping of COCO-17 coordinates to Danbooru-style pose tags. Returns tags with booru underscores (``arms_up``, ``standing``, ...) so they can be folded into the WD14 general-tag list. All checks are normalised by shoulder width and gated on keypoint confidence. For multi-person images only the primary person (highest mean keypoint confidence) is used for tag generation. Other detections are kept only for the people_count and skeleton overlay. """ if not kpts: return [] # Pick the primary person: highest mean keypoint confidence. confs = [float(np.mean(kp[:, 2])) for kp in kpts if kp.size] best = int(np.argmax(confs)) return _single_person_tags(kpts[best]) # --------------------------------------------------------------------------- # ONNX pose runner # --------------------------------------------------------------------------- def _providers() -> list[str]: avail = ort.get_available_providers() return ["CUDAExecutionProvider", "CPUExecutionProvider"] if "CUDAExecutionProvider" in avail else ["CPUExecutionProvider"] def _letterbox(im: Image.Image, size: int = 1024) -> tuple[Image.Image, float, float, float]: """Letterbox *im* to ``size × size`` with grey padding (114, 114, 114). YOLO-native preprocessing: scale so the image fits inside the box, then pad the remaining space with neutral grey. This preserves the full image content (no cropping). Returns ``(canvas, scale, pad_x, pad_y)`` so the caller can invert the mapping back to original coordinates. """ iw, ih = im.size s = min(size / float(iw), size / float(ih)) nw, nh = int(round(iw * s)), int(round(ih * s)) resized = im.resize((nw, nh), Image.LANCZOS) canvas = Image.new("RGB", (size, size), (114, 114, 114)) px = (size - nw) // 2 py = (size - nh) // 2 canvas.paste(resized, (px, py)) return canvas, s, float(px), float(py) def _inverse_letterbox(kp: np.ndarray, orig_size, pad_size: int, scale: float, pad_x: float, pad_y: float) -> np.ndarray: """Map keypoints from letterbox space back into original image coordinates. Inverts the letterbox transformation: ``(kp_x - pad_x) / scale`` for x, ``(kp_y - pad_y) / scale`` for y. Row layout is preserved: ``(N, 3)`` with the confidence column untouched. """ xs = (kp[:, 0] - pad_x) / scale ys = (kp[:, 1] - pad_y) / scale return np.stack([xs, ys, kp[:, 2]], axis=1) class PoseEstimator: """YOLO-pose wrapper with lazy model download.""" def __init__(self, repo_id: str = _POSE_REPO, filename: str = _POSE_FILENAME): self._repo = repo_id self._filename = filename self._session: Optional["ort.InferenceSession"] = None self._input_shape: tuple[int, int] = (_POSE_INPUT, _POSE_INPUT) 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, keypoints}. ``keypoints`` is a list of ``(N, 3)`` float arrays (COCO-17 order, columns x/y/confidence) mapped back into the coordinates of the EXIF-transposed input image, so callers can overlay skeletons or feed a ControlNet pipeline directly. """ if not self.ensure_loaded(): return {"pose_tags": [], "people_count": 0, "keypoints": []} pil = image if isinstance(image, Image.Image) else Image.fromarray(np.asarray(image)) # Transpose first so the fitted space and the display orientation # match; the inverse-letterbox mapping then needs no EXIF bookkeeping. pil = ImageOps.exif_transpose(pil) orig_size = pil.size size = self._input_shape[0] canvas, scale, pad_x, pad_y = _letterbox(pil, size) arr = np.asarray(canvas, 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)}") # Compute per-detection mean confidence and map keypoints back. mapped = [_inverse_letterbox(kp, orig_size, size, scale, pad_x, pad_y) for kp in kpts] confs = [float(np.mean(kp[:, 2])) for kp in kpts if kp.size] pose_score = float(np.mean(confs)) if confs else 0.0 # Quality gate: suppress tags for low-confidence detections. if pose_score >= _MIN_POSE_CONF: pose_tags = _keypoints_to_pose_tags(kpts) else: pose_tags = [] if _DEBUG_POSE: print(f"[pose] pose_score={pose_score:.3f}, tags={pose_tags}") return { "pose_tags": pose_tags, "people_count": len(kpts), "pose_score": round(pose_score, 4), "keypoints": mapped, } _pose_instance: PoseEstimator | None = None def get_pose_tagger() -> PoseEstimator: global _pose_instance if _pose_instance is None: _pose_instance = PoseEstimator() return _pose_instance