src/pose.py

"""Human pose tracking."""

from __future__ import annotations

from dataclasses import dataclass, field
from pathlib import Path

import numpy as np

from depth import _pick_device as pick_device
from smoothing import LandmarkFilter

NOSE = 0
LEFT_EYE, RIGHT_EYE = 1, 2
LEFT_EAR, RIGHT_EAR = 3, 4
LEFT_SHOULDER, RIGHT_SHOULDER = 5, 6
LEFT_ELBOW, RIGHT_ELBOW = 7, 8
LEFT_WRIST, RIGHT_WRIST = 9, 10
LEFT_HIP, RIGHT_HIP = 11, 12
LEFT_KNEE, RIGHT_KNEE = 13, 14
LEFT_ANKLE, RIGHT_ANKLE = 15, 16

KEYPOINT_NAMES: tuple[str, ...] = (
    "nose", "l eye", "r eye", "l ear", "r ear",
    "l shoulder", "r shoulder", "l elbow", "r elbow", "l wrist", "r wrist",
    "l hip", "r hip", "l knee", "r knee", "l ankle", "r ankle",
)

SKELETON: tuple[tuple[int, int], ...] = (
    (LEFT_SHOULDER, RIGHT_SHOULDER), (LEFT_SHOULDER, LEFT_HIP),
    (RIGHT_SHOULDER, RIGHT_HIP), (LEFT_HIP, RIGHT_HIP),
    (LEFT_SHOULDER, LEFT_ELBOW), (LEFT_ELBOW, LEFT_WRIST),
    (RIGHT_SHOULDER, RIGHT_ELBOW), (RIGHT_ELBOW, RIGHT_WRIST),
    (LEFT_HIP, LEFT_KNEE), (LEFT_KNEE, LEFT_ANKLE),
    (RIGHT_HIP, RIGHT_KNEE), (RIGHT_KNEE, RIGHT_ANKLE),
    (NOSE, LEFT_EYE), (NOSE, RIGHT_EYE),
    (LEFT_EYE, LEFT_EAR), (RIGHT_EYE, RIGHT_EAR),
    (LEFT_EAR, LEFT_SHOULDER), (RIGHT_EAR, RIGHT_SHOULDER),
)

FACE = (NOSE, LEFT_EYE, RIGHT_EYE, LEFT_EAR, RIGHT_EAR)


@dataclass
class Pose:
    """One detected person."""

    keypoints: np.ndarray
    scores: np.ndarray
    box: np.ndarray
    confidence: float = 0.0
    track_id: int = -1
    filter: LandmarkFilter | None = field(default=None, repr=False)

    @property
    def center(self) -> np.ndarray:
        """Box center point."""
        return np.array([(self.box[0] + self.box[2]) * 0.5,
                         (self.box[1] + self.box[3]) * 0.5])

    @property
    def height(self) -> float:
        """Bounding box height."""
        return float(max(self.box[3] - self.box[1], 1.0))

    def visible(self, index: int, threshold: float = 0.35) -> bool:
        """One keypoint is reliable."""
        return bool(self.scores[index] >= threshold)

    def point(self, *indices: int, threshold: float = 0.35) -> np.ndarray | None:
        """Mean of the reliable keypoints."""
        good = [self.keypoints[i] for i in indices if self.visible(i, threshold)]
        if not good:
            return None
        return np.mean(good, axis=0)


class PoseTracker:
    """Body keypoints from a YOLO pose model."""

    def __init__(self, model_path: str | Path, imgsz: int = 640, conf: float = 0.35,
                 max_people: int = 4, smooth: bool = True) -> None:
        from ultralytics import YOLO

        model_path = Path(model_path)
        if not model_path.exists():
            raise FileNotFoundError(f"Pose model not found: {model_path}")
        self.device = pick_device()
        self.model = YOLO(str(model_path))
        self.imgsz = imgsz
        self.conf = conf
        self.max_people = max_people
        self.smooth = smooth
        self._prev: list[Pose] = []
        self._filters: dict[int, LandmarkFilter] = {}
        self._next_id = 0

    def __call__(self, frame: np.ndarray, dt: float = 1 / 30) -> list[Pose]:
        """Detect people in one frame."""
        result = self.model.predict(frame, imgsz=self.imgsz, conf=self.conf,
                                    device=self.device, verbose=False)[0]
        if result.keypoints is None or len(result.keypoints) == 0:
            self._prev = []
            return []

        xy = result.keypoints.xy.cpu().numpy().astype(np.float64)
        scores = result.keypoints.conf
        scores = (np.ones(xy.shape[:2]) if scores is None
                  else scores.cpu().numpy().astype(np.float64))
        boxes = result.boxes.xyxy.cpu().numpy().astype(np.float64)
        confs = result.boxes.conf.cpu().numpy().astype(np.float64)

        poses = [Pose(keypoints=xy[i], scores=scores[i], box=boxes[i],
                      confidence=float(confs[i]))
                 for i in range(min(len(xy), self.max_people))]
        poses.sort(key=lambda p: p.height, reverse=True)
        self._track(poses)
        if self.smooth:
            self._smooth(poses, dt)
        self._prev = poses
        return poses

    def _track(self, poses: list[Pose]) -> None:
        """Keep ids stable between frames."""
        taken: set[int] = set()
        for p in poses:
            c = p.center
            best, best_d = None, float("inf")
            for old in self._prev:
                if old.track_id in taken:
                    continue
                d = float(np.linalg.norm(c - old.center))
                if d < best_d:
                    best, best_d = old, d
            if best is not None and best_d < max(best.height * 0.6, 60.0):
                p.track_id = best.track_id
            else:
                p.track_id = self._next_id
                self._next_id += 1
            taken.add(p.track_id)
        alive = {p.track_id for p in poses}
        for tid in [t for t in self._filters if t not in alive]:
            self._filters.pop(tid)

    def _smooth(self, poses: list[Pose], dt: float) -> None:
        """Filter keypoints per track."""
        for p in poses:
            flt = self._filters.get(p.track_id)
            if flt is None:
                flt = LandmarkFilter(min_cutoff=0.7, beta=0.02)
                self._filters[p.track_id] = flt
            p.keypoints = flt(p.keypoints, dt)
            p.filter = flt