from __future__ import annotations
from dataclasses import dataclass
from pathlib import Path
import time
import urllib.request
import cv2
import numpy as np
@dataclass
class PoseResult17:
"""Common COCO-17 pose result returned by all backends."""
xy: np.ndarray # (17, 2), pixel coordinates
conf: np.ndarray # (17,), confidence/visibility
bbox: np.ndarray | None # [x1, y1, x2, y2]
class PoseBackendBase:
name = "unknown"
def __init__(self):
self.last_inference_ms = 0.0
def infer(self, frame_bgr: np.ndarray) -> PoseResult17 | None:
raise NotImplementedError
def close(self) -> None:
pass
class UltralyticsPoseBackend(PoseBackendBase):
name = "Ultralytics"
def __init__(
self,
model_name: str = "yolo26n-pose.pt",
imgsz: int = 640,
person_conf: float = 0.30,
):
super().__init__()
import torch
from ultralytics import YOLO
self.model = YOLO(model_name)
self.imgsz = imgsz
self.person_conf = person_conf
self.use_cuda = torch.cuda.is_available()
self.device = 0 if self.use_cuda else "cpu"
print("[Pose] backend = Ultralytics")
print(f"[Pose] model = {model_name}")
print(f"[Pose] device = {self.device}")
if self.use_cuda:
print(f"[Pose] GPU = {torch.cuda.get_device_name(0)}")
def infer(self, frame_bgr: np.ndarray) -> PoseResult17 | None:
kwargs = dict(
source=frame_bgr,
imgsz=self.imgsz,
conf=self.person_conf,
device=self.device,
verbose=False,
)
# Current Ultralytics:
# quantize=16 replaces deprecated half=True.
if self.use_cuda:
kwargs["quantize"] = 16
t0 = time.perf_counter()
result = self.model.predict(**kwargs)[0]
self.last_inference_ms = (
time.perf_counter() - t0
) * 1000.0
if (
result.boxes is None
or result.keypoints is None
or len(result.boxes) == 0
):
return None
# Single-subject assumption:
# choose the largest detected person.
boxes = (
result.boxes.xyxy
.detach()
.cpu()
.numpy()
)
areas = (
(boxes[:, 2] - boxes[:, 0])
* (boxes[:, 3] - boxes[:, 1])
)
index = int(np.argmax(areas))
xy = (
result.keypoints.xy[index]
.detach()
.cpu()
.numpy()
.astype(np.float32)
)
conf_tensor = result.keypoints.conf
if conf_tensor is None:
conf = np.ones(
17,
dtype=np.float32
)
else:
conf = (
conf_tensor[index]
.detach()
.cpu()
.numpy()
.astype(np.float32)
)
bbox = boxes[index].astype(
np.float32
)
return PoseResult17(
xy=xy,
conf=conf,
bbox=bbox
)
# ============================================================
# MediaPipe Pose Landmarker 33 -> COCO 17
# ============================================================
MEDIAPIPE_TO_COCO17 = [
0, # Nose
2, # Left eye
5, # Right eye
7, # Left ear
8, # Right ear
11, # Left shoulder
12, # Right shoulder
13, # Left elbow
14, # Right elbow
15, # Left wrist
16, # Right wrist
23, # Left hip
24, # Right hip
25, # Left knee
26, # Right knee
27, # Left ankle
28, # Right ankle
]
MEDIAPIPE_MODELS = {
"lite": (
"pose_landmarker_lite.task",
"https://storage.googleapis.com/"
"mediapipe-models/pose_landmarker/"
"pose_landmarker_lite/float16/1/"
"pose_landmarker_lite.task",
),
"full": (
"pose_landmarker_full.task",
"https://storage.googleapis.com/"
"mediapipe-models/pose_landmarker/"
"pose_landmarker_full/float16/1/"
"pose_landmarker_full.task",
),
"heavy": (
"pose_landmarker_heavy.task",
"https://storage.googleapis.com/"
"mediapipe-models/pose_landmarker/"
"pose_landmarker_heavy/float16/1/"
"pose_landmarker_heavy.task",
),
}
class MediaPipePoseBackend(PoseBackendBase):
name = "MediaPipe"
def __init__(
self,
model_variant: str = "lite",
detection_conf: float = 0.40,
presence_conf: float = 0.40,
tracking_conf: float = 0.40,
model_dir: str = "models",
):
super().__init__()
import mediapipe as mp
if model_variant not in MEDIAPIPE_MODELS:
raise ValueError(
"model_variant must be "
"lite, full, or heavy"
)
self.mp = mp
self.model_variant = model_variant
self.last_timestamp_ms = -1
filename, url = (
MEDIAPIPE_MODELS[
model_variant
]
)
model_dir_path = Path(
model_dir
)
model_dir_path.mkdir(
parents=True,
exist_ok=True
)
model_path = (
model_dir_path / filename
)
# ----------------------------------------------------
# Automatic model download
# ----------------------------------------------------
if not model_path.exists():
print(
"[Pose] downloading "
"MediaPipe model:"
)
print(
f" {url}"
)
urllib.request.urlretrieve(
url,
model_path
)
print(
f"[Pose] saved to "
f"{model_path}"
)
# ----------------------------------------------------
# VIDEO mode:
# tracking information is reused between frames.
# ----------------------------------------------------
options = (
mp.tasks.vision
.PoseLandmarkerOptions(
base_options=(
mp.tasks.BaseOptions(
model_asset_path=
str(model_path)
)
),
running_mode=(
mp.tasks.vision
.RunningMode.VIDEO
),
num_poses=1,
min_pose_detection_confidence=
detection_conf,
min_pose_presence_confidence=
presence_conf,
min_tracking_confidence=
tracking_conf,
output_segmentation_masks=False,
)
)
self.landmarker = (
mp.tasks.vision
.PoseLandmarker
.create_from_options(
options
)
)
print(
"[Pose] backend = MediaPipe"
)
print(
f"[Pose] model = "
f"{model_variant}"
)
print(
"[Pose] mode = VIDEO"
)
@staticmethod
def _landmark_confidence(lm):
values = []
visibility = getattr(
lm,
"visibility",
None
)
presence = getattr(
lm,
"presence",
None
)
if visibility is not None:
values.append(
float(visibility)
)
if presence is not None:
values.append(
float(presence)
)
if values:
return min(values)
return 1.0
def infer(
self,
frame_bgr: np.ndarray
) -> PoseResult17 | None:
h, w = frame_bgr.shape[:2]
# BGR -> RGB
frame_rgb = cv2.cvtColor(
frame_bgr,
cv2.COLOR_BGR2RGB
)
frame_rgb = np.ascontiguousarray(
frame_rgb
)
mp_image = self.mp.Image(
image_format=
self.mp.ImageFormat.SRGB,
data=frame_rgb,
)
# VIDEO mode requires strictly increasing timestamps.
timestamp_ms = int(
time.perf_counter()
* 1000.0
)
timestamp_ms = max(
timestamp_ms,
self.last_timestamp_ms + 1
)
self.last_timestamp_ms = (
timestamp_ms
)
t0 = time.perf_counter()
result = (
self.landmarker
.detect_for_video(
mp_image,
timestamp_ms,
)
)
self.last_inference_ms = (
time.perf_counter() - t0
) * 1000.0
if not result.pose_landmarks:
return None
landmarks = (
result.pose_landmarks[0]
)
xy = np.zeros(
(17, 2),
dtype=np.float32
)
conf = np.zeros(
17,
dtype=np.float32
)
# ----------------------------------------------------
# MediaPipe 33 -> COCO17
# ----------------------------------------------------
for coco_id, mp_id in enumerate(
MEDIAPIPE_TO_COCO17
):
lm = landmarks[mp_id]
xy[coco_id, 0] = (
float(lm.x) * w
)
xy[coco_id, 1] = (
float(lm.y) * h
)
conf[coco_id] = (
self._landmark_confidence(
lm
)
)
# ----------------------------------------------------
# Approximate bbox from landmarks
# ----------------------------------------------------
good = conf >= 0.20
bbox = None
if np.count_nonzero(good) >= 4:
pts = xy[good]
pad = 15.0
bbox = np.array(
[
max(
0.0,
float(
pts[:, 0].min()
- pad
)
),
max(
0.0,
float(
pts[:, 1].min()
- pad
)
),
min(
float(w - 1),
float(
pts[:, 0].max()
+ pad
)
),
min(
float(h - 1),
float(
pts[:, 1].max()
+ pad
)
),
],
dtype=np.float32,
)
return PoseResult17(
xy=xy,
conf=conf,
bbox=bbox
)
def close(self):
self.landmarker.close()
# ============================================================
# Factory
# ============================================================
def create_pose_backend(
backend: str,
*,
ultralytics_model:
str = "yolo26n-pose.pt",
ultralytics_imgsz:
int = 640,
person_conf:
float = 0.30,
mediapipe_model:
str = "lite",
):
backend = backend.lower()
if backend == "ultralytics":
return UltralyticsPoseBackend(
model_name=
ultralytics_model,
imgsz=
ultralytics_imgsz,
person_conf=
person_conf,
)
if backend == "mediapipe":
return MediaPipePoseBackend(
model_variant=
mediapipe_model,
)
raise ValueError(
f"Unknown backend: {backend}"
)
