File size: 5,458 Bytes
9f85448 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 | """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
|