File size: 2,589 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 | from __future__ import annotations
import os
import time
from pathlib import Path
os.environ.setdefault("GLOG_minloglevel", "3")
import numpy as np
from hand import Hand
class MediaPipeTracker:
"""MediaPipe hand landmarker."""
def __init__(
self,
model_path: str | Path,
conf: float = 0.4,
max_hands: int = 2,
video_mode: bool = True,
**_ignored,
) -> None:
import mediapipe as mp
from mediapipe.tasks.python import BaseOptions, vision
model_path = Path(model_path)
if not model_path.exists():
raise FileNotFoundError(f"Model not found: {model_path}")
self._mp = mp
self.device = "cpu"
self.max_hands = max_hands
self.video_mode = video_mode
mode = vision.RunningMode.VIDEO if video_mode else vision.RunningMode.IMAGE
options = vision.HandLandmarkerOptions(
base_options=BaseOptions(model_asset_path=str(model_path)),
running_mode=mode,
num_hands=max_hands,
min_hand_detection_confidence=conf,
min_hand_presence_confidence=conf,
min_tracking_confidence=0.5,
)
self.detector = vision.HandLandmarker.create_from_options(options)
self._t0 = time.perf_counter()
self._last_ts = -1
def __call__(self, frame: np.ndarray) -> list[Hand]:
"""Detect hands in frame."""
h, w = frame.shape[:2]
rgb = frame[:, :, ::-1].copy()
image = self._mp.Image(image_format=self._mp.ImageFormat.SRGB, data=rgb)
if self.video_mode:
ts = max(int((time.perf_counter() - self._t0) * 1000), self._last_ts + 1)
self._last_ts = ts
result = self.detector.detect_for_video(image, ts)
else:
result = self.detector.detect(image)
hands: list[Hand] = []
for i, lms in enumerate(result.hand_landmarks):
kp = np.array([[lm.x * w, lm.y * h] for lm in lms], dtype=np.float64)
conf = 1.0
if result.handedness and i < len(result.handedness):
conf = float(result.handedness[i][0].score)
box = np.array([kp[:, 0].min(), kp[:, 1].min(),
kp[:, 0].max(), kp[:, 1].max()], dtype=np.float64)
hands.append(Hand(keypoints=kp, scores=np.ones(21), box=box, conf=conf))
hands.sort(key=lambda x: x.conf, reverse=True)
return hands[: self.max_hands]
def close(self) -> None:
"""Release detector."""
self.detector.close()
|