| """ |
| MediaPipe landmark extraction (pose + left hand + right hand). |
| |
| Two backends are supported and selected automatically: |
| |
| 1. **Tasks API** (``mediapipe >= 0.10.18`` and ``mediapipe 1.x``) using the |
| ``holistic_landmarker.task`` bundle. This is the only backend available on |
| MediaPipe 1.0.1 / Python 3.13, which is what this project runs on |
| (``mediapipe.solutions`` was removed in 1.x). |
| 2. **Legacy Solutions API** (``mediapipe.solutions.holistic``) when present. |
| |
| Both produce exactly the same :class:`~backend.preprocessing.landmarks.FrameLandmarks` |
| object, so the rest of the system is backend-agnostic. |
| |
| Usage |
| ----- |
| from backend.preprocessing.mediapipe_extractor import MediaPipeExtractor |
| |
| with MediaPipeExtractor() as ex: |
| frame = ex.extract_frame(rgb_uint8) # FrameLandmarks |
| seq, coords, stats = ex.extract_video("clip.mp4") |
| """ |
|
|
| from __future__ import annotations |
|
|
| import logging |
| from dataclasses import dataclass, field |
| from pathlib import Path |
| from typing import Dict, Iterable, List, Optional, Tuple |
|
|
| import numpy as np |
|
|
| import config |
| from .landmarks import FrameLandmarks, pack_frame |
|
|
| logger = logging.getLogger(__name__) |
|
|
| DEFAULT_MODEL_ASSET = config.MODELS_DIR / "mediapipe_assets" / "holistic_landmarker.task" |
|
|
| |
| N_HAND = 21 |
| N_POSE = 33 |
| N_FACE = 468 |
|
|
|
|
| def _to_xyz(landmarks: Iterable) -> Tuple[np.ndarray, np.ndarray]: |
| """Convert a MediaPipe landmark list into (N,3) coords and (N,) visibility.""" |
| pts = list(landmarks) |
| coords = np.array([[float(p.x), float(p.y), float(p.z)] for p in pts], dtype=np.float32) |
| vis = np.array([float(getattr(p, "visibility", 1.0) or 0.0) for p in pts], dtype=np.float32) |
| return coords, vis |
|
|
|
|
| @dataclass |
| class ExtractionStats: |
| """Counters describing how well MediaPipe did on a video.""" |
|
|
| video: str = "" |
| frames_read: int = 0 |
| frames_decode_error: int = 0 |
| frames_with_pose: int = 0 |
| frames_with_left_hand: int = 0 |
| frames_with_right_hand: int = 0 |
| frames_with_any_hand: int = 0 |
| frames_with_nothing: int = 0 |
| source_fps: float = 0.0 |
| width: int = 0 |
| height: int = 0 |
| duration_s: float = 0.0 |
| error: str = "" |
|
|
| @property |
| def extraction_success(self) -> float: |
| """Fraction of decoded frames with at least one usable landmark group.""" |
| usable = self.frames_read - self.frames_with_nothing |
| return float(usable) / float(self.frames_read) if self.frames_read else 0.0 |
|
|
| @property |
| def both_hands_rate(self) -> float: |
| both = 0 if not self.frames_read else min(self.frames_with_left_hand, self.frames_with_right_hand) |
| return float(both) / float(self.frames_read) if self.frames_read else 0.0 |
|
|
| def as_dict(self) -> Dict[str, float]: |
| d = self.__dict__.copy() |
| d["extraction_success"] = self.extraction_success |
| d["both_hands_rate"] = self.both_hands_rate |
| return d |
|
|
|
|
| class MediaPipeExtractor: |
| """Thin, defensive wrapper around MediaPipe holistic landmark extraction.""" |
|
|
| def __init__( |
| self, |
| model_asset_path: Optional[Path | str] = None, |
| backend: str = "auto", |
| min_detection_confidence: float = 0.5, |
| min_tracking_confidence: float = 0.5, |
| min_hand_confidence: float = 0.5, |
| delegate: str = "cpu", |
| ) -> None: |
| self.model_asset_path = Path(model_asset_path or DEFAULT_MODEL_ASSET) |
| self.backend = backend |
| self.min_detection_confidence = min_detection_confidence |
| self.min_tracking_confidence = min_tracking_confidence |
| self.min_hand_confidence = min_hand_confidence |
| self.delegate = delegate |
| self._impl = None |
| self._mode = None |
| self._timestamp_ms = 0 |
| self._frame_shape: Optional[Tuple[int, int]] = None |
| self._reinit_count = 0 |
| self._init_backend() |
|
|
| |
| |
| |
| def _init_backend(self) -> None: |
| errors: List[str] = [] |
| if self.backend in ("auto", "tasks"): |
| try: |
| self._init_tasks() |
| self._mode = "tasks" |
| logger.info("MediaPipe backend: Tasks API (holistic_landmarker.task)") |
| return |
| except Exception as exc: |
| errors.append(f"tasks: {exc}") |
| if self.backend in ("auto", "legacy"): |
| try: |
| self._init_legacy() |
| self._mode = "legacy" |
| logger.info("MediaPipe backend: legacy solutions.holistic") |
| return |
| except Exception as exc: |
| errors.append(f"legacy: {exc}") |
| raise RuntimeError( |
| "Could not initialise a MediaPipe backend.\n " |
| + "\n ".join(errors) |
| + f"\n Tasks API needs the model bundle: {self.model_asset_path}\n" |
| " Run: python scripts/download_models.py" |
| ) |
|
|
| def _init_tasks(self) -> None: |
| import mediapipe as mp |
| from mediapipe.tasks import python as mp_python |
| from mediapipe.tasks.python import vision |
|
|
| if not self.model_asset_path.exists(): |
| raise FileNotFoundError(f"model bundle not found: {self.model_asset_path}") |
| base_options = mp_python.BaseOptions( |
| model_asset_path=str(self.model_asset_path), |
| delegate=mp_python.BaseOptions.Delegate.GPU |
| if self.delegate == "gpu" |
| else mp_python.BaseOptions.Delegate.CPU, |
| ) |
| options = vision.HolisticLandmarkerOptions( |
| base_options=base_options, |
| running_mode=vision.RunningMode.VIDEO, |
| min_face_detection_confidence=self.min_detection_confidence, |
| min_face_landmarks_confidence=self.min_detection_confidence, |
| min_pose_detection_confidence=self.min_detection_confidence, |
| min_pose_landmarks_confidence=self.min_detection_confidence, |
| min_hand_landmarks_confidence=self.min_hand_confidence, |
| output_face_blendshapes=False, |
| output_segmentation_mask=False, |
| ) |
| self._vision = vision |
| self._mp = mp |
| self._impl = vision.HolisticLandmarker.create_from_options(options) |
|
|
| def _init_legacy(self) -> None: |
| import mediapipe as mp |
|
|
| if not hasattr(mp, "solutions"): |
| raise AttributeError("mediapipe.solutions not available in this build") |
| self._mp = mp |
| self._impl = mp.solutions.holistic.Holistic( |
| static_image_mode=False, |
| model_complexity=1, |
| min_detection_confidence=self.min_detection_confidence, |
| min_tracking_confidence=self.min_tracking_confidence, |
| ) |
|
|
| def _ensure_frame_shape(self, shape: Tuple[int, int]) -> None: |
| """The Tasks API video graph requires a constant frame size. |
| |
| Switching camera resolution, or feeding a new video with different |
| dimensions, otherwise kills the graph with an INTERNAL error |
| (``SegmentationSmoothingCalculator ... current_mat->cols == previous_mat->cols``). |
| We transparently rebuild the landmarker when the size changes. |
| """ |
| if self._mode != "tasks" or self._frame_shape is None: |
| self._frame_shape = shape |
| return |
| if shape != self._frame_shape: |
| logger.debug("frame size changed %s -> %s; rebuilding landmarker", |
| self._frame_shape, shape) |
| self.close() |
| self._init_tasks() |
| self._frame_shape = shape |
| self._reinit_count += 1 |
|
|
| |
| |
| |
| def extract_frame(self, rgb: np.ndarray, timestamp_ms: Optional[int] = None) -> FrameLandmarks: |
| """Extract pose + both hands from one RGB uint8 frame.""" |
| if rgb is None or rgb.size == 0: |
| return FrameLandmarks() |
| rgb = np.ascontiguousarray(rgb[:, :, ::-1] if rgb.shape[-1] == 3 and rgb.dtype == np.float32 else rgb) |
| self._ensure_frame_shape((int(rgb.shape[0]), int(rgb.shape[1]))) |
|
|
| if self._mode == "tasks": |
| image = self._mp.Image(image_format=self._mp.ImageFormat.SRGB, |
| data=np.ascontiguousarray(rgb.astype(np.uint8))) |
| |
| ts = int(timestamp_ms) if timestamp_ms is not None else self._timestamp_ms |
| ts = max(ts, self._timestamp_ms) |
| self._timestamp_ms = ts + 1 |
| res = self._impl.detect_for_video(image, ts) |
| pose = getattr(res, "pose_landmarks", None) |
| lh = getattr(res, "left_hand_landmarks", None) |
| rh = getattr(res, "right_hand_landmarks", None) |
| pose_c, pose_v = _to_xyz(pose) if pose else (None, None) |
| lh_c = _to_xyz(lh)[0] if lh else None |
| rh_c = _to_xyz(rh)[0] if rh else None |
| else: |
| res = self._impl.process(np.ascontiguousarray(rgb)) |
| pose_c, pose_v = _to_xyz(res.pose_landmarks.landmark) if res.pose_landmarks else (None, None) |
| lh_c = _to_xyz(res.left_hand_landmarks.landmark)[0] if res.left_hand_landmarks else None |
| rh_c = _to_xyz(res.right_hand_landmarks.landmark)[0] if res.right_hand_landmarks else None |
|
|
| return FrameLandmarks(pose=pose_c, left_hand=lh_c, right_hand=rh_c, pose_visibility=pose_v) |
|
|
| |
| |
| |
| def extract_video( |
| self, |
| video_path: Path | str, |
| max_frames: Optional[int] = None, |
| frame_stride: int = 1, |
| sample_every: Optional[int] = None, |
| ) -> Tuple[np.ndarray, np.ndarray, ExtractionStats]: |
| """Extract landmarks from a video file. |
| |
| Parameters |
| ---------- |
| video_path : path to a video readable by OpenCV |
| max_frames : stop after this many *decoded* frames (privacy/speed guard) |
| frame_stride : decode every n-th frame (``sample_every`` is an alias) |
| |
| Returns |
| ------- |
| coords : (T, 75, 3) float32 in normalised image coordinates |
| mask : (T, 75) float32 presence mask |
| stats : :class:`ExtractionStats` |
| """ |
| import cv2 |
|
|
| video_path = Path(video_path) |
| stride = int(sample_every or frame_stride) |
| stats = ExtractionStats(video=str(video_path)) |
| coords_out: List[np.ndarray] = [] |
| masks_out: List[np.ndarray] = [] |
|
|
| cap = cv2.VideoCapture(str(video_path)) |
| if not cap.isOpened(): |
| stats.error = "video_open_failed" |
| return (np.zeros((0, config.NUM_LANDMARKS, config.COORDS_PER_LANDMARK), dtype=np.float32), |
| np.zeros((0, config.NUM_LANDMARKS), dtype=np.float32), stats) |
|
|
| stats.source_fps = float(cap.get(cv2.CAP_PROP_FPS) or 0.0) |
| stats.width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH) or 0) |
| stats.height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT) or 0) |
| total = float(cap.get(cv2.CAP_PROP_FRAME_COUNT) or 0.0) |
| if stats.source_fps > 0: |
| stats.duration_s = total / stats.source_fps |
|
|
| limit = int(max_frames or config.MAX_FRAMES_READ) |
| idx = 0 |
| try: |
| while idx < limit: |
| ok, frame_bgr = cap.read() |
| if not ok: |
| break |
| idx += 1 |
| if stride > 1 and (idx - 1) % stride != 0: |
| continue |
| stats.frames_read += 1 |
| rgb = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB) |
| ts = int(1000.0 * (idx - 1) / stats.source_fps) if stats.source_fps > 0 else idx * 33 |
| lm = self.extract_frame(rgb, timestamp_ms=ts) |
| c, m = pack_frame(lm) |
| coords_out.append(c) |
| masks_out.append(m) |
|
|
| stats.frames_with_pose += int(lm.pose is not None) |
| stats.frames_with_left_hand += int(lm.left_hand is not None) |
| stats.frames_with_right_hand += int(lm.right_hand is not None) |
| stats.frames_with_any_hand += int(lm.left_hand is not None or lm.right_hand is not None) |
| if lm.is_empty(): |
| stats.frames_with_nothing += 1 |
| except Exception as exc: |
| stats.error = f"{type(exc).__name__}: {exc}" |
| finally: |
| cap.release() |
|
|
| if not coords_out: |
| return (np.zeros((0, config.NUM_LANDMARKS, config.COORDS_PER_LANDMARK), dtype=np.float32), |
| np.zeros((0, config.NUM_LANDMARKS), dtype=np.float32), stats) |
| return (np.stack(coords_out).astype(np.float32), np.stack(masks_out).astype(np.float32), stats) |
|
|
| def reset(self) -> None: |
| """Reset tracking state between videos (keeps the graph alive).""" |
| self._timestamp_ms = 0 |
|
|
| def close(self) -> None: |
| try: |
| if self._impl is not None and hasattr(self._impl, "close"): |
| self._impl.close() |
| except Exception: |
| pass |
| self._impl = None |
|
|
| def __enter__(self) -> "MediaPipeExtractor": |
| return self |
|
|
| def __exit__(self, *exc) -> None: |
| self.close() |
|
|
| def __del__(self): |
| self.close() |
|
|
|
|
| def extract_landmarks_for_frames( |
| frames: Iterable[np.ndarray], extractor: Optional[MediaPipeExtractor] = None |
| ) -> Tuple[np.ndarray, np.ndarray]: |
| """Convenience helper: run extraction over an iterable of RGB frames.""" |
| own = extractor is None |
| ex = extractor or MediaPipeExtractor() |
| coords, masks = [], [] |
| for i, f in enumerate(frames): |
| c, m = pack_frame(ex.extract_frame(f, timestamp_ms=i * 33)) |
| coords.append(c) |
| masks.append(m) |
| if own: |
| ex.close() |
| if not coords: |
| return (np.zeros((0, config.NUM_LANDMARKS, config.COORDS_PER_LANDMARK), dtype=np.float32), |
| np.zeros((0, config.NUM_LANDMARKS), dtype=np.float32)) |
| return np.stack(coords), np.stack(masks) |
|
|
