63 lines
2.2 KiB
Python
63 lines
2.2 KiB
Python
from __future__ import annotations
|
|
|
|
import threading
|
|
import time
|
|
from typing import Callable
|
|
|
|
import mediapipe as mp
|
|
from loguru import logger
|
|
|
|
from app.vision.pose_models import DEFAULT_MODEL_PATH
|
|
|
|
PoseLandmarker = mp.tasks.vision.PoseLandmarker
|
|
PoseLandmarkerOptions = mp.tasks.vision.PoseLandmarkerOptions
|
|
VisionRunningMode = mp.tasks.vision.RunningMode
|
|
BaseOptions = mp.tasks.BaseOptions
|
|
|
|
class PoseLandmarkerWrapper:
|
|
"""MediaPipe姿态关键点检测器封装"""
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
model_path: str | None = None,
|
|
prefer_gpu: bool = True,
|
|
result_callback: Callable | None = None,
|
|
) -> None:
|
|
"""初始化姿态检测器，优先尝试GPU委托，失败则回退到CPU"""
|
|
self.model_path = model_path or DEFAULT_MODEL_PATH
|
|
|
|
if prefer_gpu:
|
|
try:
|
|
self.delegate = BaseOptions.Delegate.GPU
|
|
self._landmarker = self._create(PoseLandmarker.Delegate.GPU)
|
|
logger.info("MediaPipe PoseLandmarker initialized with GPU delegate")
|
|
return
|
|
except Exception as exc:
|
|
logger.warning("MediaPipe GPU delegate unavailable, falling back to CPU: {}", exc)
|
|
|
|
self.delegate = BaseOptions.Delegate.CPU
|
|
self._landmarker = self._create(PoseLandmarker.Delegate.CPU, result_callback)
|
|
logger.info("MediaPipe PoseLandmarker initialized with CPU delegate")
|
|
|
|
def _create(self, delegate, result_callback=None):
|
|
"""根据委托类型和回调创建PoseLandmarker实例"""
|
|
options = PoseLandmarkerOptions(
|
|
base_options=BaseOptions(model_asset_path=self.model_path, delegate=delegate),
|
|
running_mode=VisionRunningMode.LIVE_STREAM,
|
|
result_callback=result_callback,
|
|
num_poses=1,
|
|
min_pose_detection_confidence=0.5,
|
|
min_pose_presence_confidence=0.5,
|
|
min_tracking_confidence=0.5,
|
|
)
|
|
return PoseLandmarker.create_from_options(options)
|
|
|
|
def detect_async(self, mp_image, timestamp_ms: int) -> None:
|
|
"""异步执行姿态检测"""
|
|
return self._landmarker.detect_async(mp_image, timestamp_ms)
|
|
|
|
def close(self) -> None:
|
|
"""释放MediaPipe资源"""
|
|
self._landmarker.close()
|