src/pose_tagger.py

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