kinemotion 0.10.6__py3-none-any.whl → 0.67.0__py3-none-any.whl

kinemotion/core/pose.py

@@ -1,84 +1,208 @@
-"""Pose tracking using MediaPipe Pose."""
+"""Pose tracking using MediaPipe Tasks API.
+
+The MediaPipe Solutions API was removed in version 0.10.31.
+This module now uses the Tasks API (PoseLandmarker).
+
+Key differences from Solution API:
+- Tasks API uses index-based landmark access (0-32) instead of enums
+- Running modes: IMAGE, VIDEO, LIVE_STREAM
+- No smooth_landmarks option (built into VIDEO mode)
+- Has min_pose_presence_confidence parameter (no Solution API equivalent)
+
+Configuration strategies for matching Solution API behavior:
+- "video": Standard VIDEO mode with temporal smoothing
+- "video_low_presence": VIDEO mode with lower min_pose_presence_confidence (0.2)
+- "video_very_low_presence": VIDEO mode with very low min_pose_presence_confidence (0.1)
+- "image": IMAGE mode (no temporal smoothing, relies on our smoothing)
+"""
+
+from __future__ import annotations
 
 import cv2
 import mediapipe as mp
 import numpy as np
 
+from .pose_landmarks import KINEMOTION_LANDMARKS, LANDMARK_INDICES
+from .timing import NULL_TIMER, Timer
+
+# Running modes
+_RUNNING_MODES = {
+    "image": mp.tasks.vision.RunningMode.IMAGE,  # type: ignore[attr-defined]
+    "video": mp.tasks.vision.RunningMode.VIDEO,  # type: ignore[attr-defined]
+}
+
+# Strategy configurations
+_STRATEGY_CONFIGS: dict[str, dict[str, float | str]] = {
+    "video": {
+        "min_pose_presence_confidence": 0.5,
+        "running_mode": "video",
+    },
+    "video_low_presence": {
+        "min_pose_presence_confidence": 0.2,
+        "running_mode": "video",
+    },
+    "video_very_low_presence": {
+        "min_pose_presence_confidence": 0.1,
+        "running_mode": "video",
+    },
+    "image": {
+        "min_pose_presence_confidence": 0.5,
+        "running_mode": "image",
+    },
+}
+
 
 class PoseTracker:
-    """Tracks human pose landmarks in video frames using MediaPipe."""
+    """Tracks human pose landmarks in video frames using MediaPipe Tasks API.
+
+    Args:
+        min_detection_confidence: Minimum confidence for pose detection (0.0-1.0)
+        min_tracking_confidence: Minimum confidence for pose tracking (0.0-1.0)
+        model_type: Model variant ("lite", "full", "heavy")
+        strategy: Configuration strategy ("video", "video_low_presence", "image")
+        timer: Optional Timer for measuring operations
+
+    Note: The Solution API's smooth_landmarks parameter cannot be replicated
+    exactly. VIDEO mode has built-in temporal smoothing that cannot be disabled.
+    """
 
-    def __init__(
+    def __init__(  # noqa: PLR0913
         self,
         min_detection_confidence: float = 0.5,
         min_tracking_confidence: float = 0.5,
-    ):
-        """
-        Initialize the pose tracker.
+        model_type: str = "lite",
+        strategy: str = "video_low_presence",
+        timer: Timer | None = None,
+    ) -> None:
+        """Initialize the pose tracker."""
+        self.timer = timer or NULL_TIMER
+        self.mp_pose = mp.tasks.vision  # type: ignore[attr-defined]
+        self.model_type = model_type
+        self.strategy = strategy
 
-        Args:
-            min_detection_confidence: Minimum confidence for pose detection
-            min_tracking_confidence: Minimum confidence for pose tracking
-        """
-        self.mp_pose = mp.solutions.pose
-        self.pose = self.mp_pose.Pose(
-            min_detection_confidence=min_detection_confidence,
+        # Get strategy configuration
+        config = _STRATEGY_CONFIGS.get(strategy, _STRATEGY_CONFIGS["video_low_presence"])
+        min_pose_presence = config["min_pose_presence_confidence"]
+        running_mode_name = str(config["running_mode"])
+        running_mode = _RUNNING_MODES[running_mode_name]
+
+        # Get model path
+        from .model_downloader import get_model_path
+
+        model_path = str(get_model_path(model_type))
+
+        # Create base options
+        base_options = mp.tasks.BaseOptions(model_asset_path=model_path)  # type: ignore[attr-defined]
+
+        # Create pose landmarker options
+        options = mp.tasks.vision.PoseLandmarkerOptions(  # type: ignore[attr-defined]
+            base_options=base_options,
+            running_mode=running_mode,
+            min_pose_detection_confidence=min_detection_confidence,
+            min_pose_presence_confidence=min_pose_presence,
             min_tracking_confidence=min_tracking_confidence,
-            model_complexity=1,
+            output_segmentation_masks=False,
         )
 
+        # Create the landmarker
+        with self.timer.measure("model_load"):
+            self.landmarker = self.mp_pose.PoseLandmarker.create_from_options(options)
+
+        self.running_mode = running_mode
+
     def process_frame(
-        self, frame: np.ndarray
+        self,
+        frame: np.ndarray,
+        timestamp_ms: int = 0,
     ) -> dict[str, tuple[float, float, float]] | None:
-        """
-        Process a single frame and extract pose landmarks.
+        """Process a single frame and extract pose landmarks.
 
         Args:
             frame: BGR image frame
+            timestamp_ms: Frame timestamp in milliseconds (required for VIDEO mode)
 
         Returns:
             Dictionary mapping landmark names to (x, y, visibility) tuples,
             or None if no pose detected. Coordinates are normalized (0-1).
         """
         # Convert BGR to RGB
-        rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+        with self.timer.measure("frame_conversion"):
+            rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+
+        # Create MediaPipe Image
+        with self.timer.measure("image_creation"):
+            mp_image = mp.Image(image_format=mp.ImageFormat.SRGB, data=rgb_frame)  # type: ignore[attr-defined]
 
         # Process the frame
-        results = self.pose.process(rgb_frame)
+        with self.timer.measure("mediapipe_inference"):
+            if self.running_mode == mp.tasks.vision.RunningMode.VIDEO:  # type: ignore[attr-defined]
+                results = self.landmarker.detect_for_video(mp_image, timestamp_ms)
+            else:  # IMAGE mode
+                results = self.landmarker.detect(mp_image)
 
         if not results.pose_landmarks:
             return None
 
-        # Extract key landmarks for feet tracking and CoM estimation
-        landmarks = {}
-        landmark_names = {
-            # Feet landmarks
-            self.mp_pose.PoseLandmark.LEFT_ANKLE: "left_ankle",
-            self.mp_pose.PoseLandmark.RIGHT_ANKLE: "right_ankle",
-            self.mp_pose.PoseLandmark.LEFT_HEEL: "left_heel",
-            self.mp_pose.PoseLandmark.RIGHT_HEEL: "right_heel",
-            self.mp_pose.PoseLandmark.LEFT_FOOT_INDEX: "left_foot_index",
-            self.mp_pose.PoseLandmark.RIGHT_FOOT_INDEX: "right_foot_index",
-            # Torso landmarks for CoM estimation
-            self.mp_pose.PoseLandmark.LEFT_HIP: "left_hip",
-            self.mp_pose.PoseLandmark.RIGHT_HIP: "right_hip",
-            self.mp_pose.PoseLandmark.LEFT_SHOULDER: "left_shoulder",
-            self.mp_pose.PoseLandmark.RIGHT_SHOULDER: "right_shoulder",
-            # Additional landmarks for better CoM estimation
-            self.mp_pose.PoseLandmark.NOSE: "nose",
-            self.mp_pose.PoseLandmark.LEFT_KNEE: "left_knee",
-            self.mp_pose.PoseLandmark.RIGHT_KNEE: "right_knee",
-        }
-
-        for landmark_id, name in landmark_names.items():
-            lm = results.pose_landmarks.landmark[landmark_id]
-            landmarks[name] = (lm.x, lm.y, lm.visibility)
+        # Extract landmarks (first pose only)
+        with self.timer.measure("landmark_extraction"):
+            landmarks = _extract_landmarks_from_results(results.pose_landmarks[0])
 
         return landmarks
 
     def close(self) -> None:
-        """Release resources."""
-        self.pose.close()
+        """Release resources.
+
+        Note: Tasks API landmarker doesn't have explicit close method.
+        Resources are released when the object is garbage collected.
+        """
+        pass
+
+
+def _extract_landmarks_from_results(
+    pose_landmarks: mp.tasks.vision.components.containers.NormalizedLandmark,  # type: ignore[valid-type]
+) -> dict[str, tuple[float, float, float]]:
+    """Extract kinemotion landmarks from pose landmarker result.
+
+    Args:
+        pose_landmarks: MediaPipe pose landmarks (list of 33 landmarks)
+
+    Returns:
+        Dictionary mapping landmark names to (x, y, visibility) tuples
+    """
+    landmarks: dict[str, tuple[float, float, float]] = {}
+
+    for name in KINEMOTION_LANDMARKS:
+        idx = LANDMARK_INDICES[name]
+        if idx < len(pose_landmarks):
+            lm = pose_landmarks[idx]
+            # Tasks API uses presence in addition to visibility
+            # Use visibility for consistency with Solution API
+            visibility = getattr(lm, "visibility", 1.0)
+            landmarks[name] = (lm.x, lm.y, visibility)
+
+    return landmarks
+
+
+# Legacy compatibility aliases for Solution API enum values
+class _LegacyPoseLandmark:
+    """Compatibility shim for Solution API enum values."""
+
+    LEFT_ANKLE = 27
+    RIGHT_ANKLE = 28
+    LEFT_HEEL = 29
+    RIGHT_HEEL = 30
+    LEFT_FOOT_INDEX = 31
+    RIGHT_FOOT_INDEX = 32
+    LEFT_HIP = 23
+    RIGHT_HIP = 24
+    LEFT_SHOULDER = 11
+    RIGHT_SHOULDER = 12
+    NOSE = 0
+    LEFT_KNEE = 25
+    RIGHT_KNEE = 26
+
+
+PoseLandmark = _LegacyPoseLandmark
 
 
 def compute_center_of_mass(
@@ -106,115 +230,148 @@ def compute_center_of_mass(
         (x, y, visibility) tuple for estimated CoM position
         visibility = average visibility of all segments used
     """
-    # Define segment representatives and their weights (as fraction of body mass)
-    # Each segment uses midpoint or average of its bounding landmarks
-    segments = []
-    segment_weights = []
-    visibilities = []
+    segments: list = []
+    weights: list = []
+    visibilities: list = []
 
-    # Head segment: 8% (use nose as proxy)
+    # Add body segments
+    _add_head_segment(segments, weights, visibilities, landmarks, visibility_threshold)
+    _add_trunk_segment(segments, weights, visibilities, landmarks, visibility_threshold)
+
+    # Add bilateral limb segments
+    for side in ["left", "right"]:
+        _add_limb_segment(
+            segments,
+            weights,
+            visibilities,
+            landmarks,
+            side,
+            "hip",
+            "knee",
+            0.10,
+            visibility_threshold,
+        )
+        _add_limb_segment(
+            segments,
+            weights,
+            visibilities,
+            landmarks,
+            side,
+            "knee",
+            "ankle",
+            0.05,
+            visibility_threshold,
+        )
+        _add_foot_segment(segments, weights, visibilities, landmarks, side, visibility_threshold)
+
+    # Fallback if no segments found
+    if not segments:
+        if "left_hip" in landmarks and "right_hip" in landmarks:
+            lh_x, lh_y, lh_vis = landmarks["left_hip"]
+            rh_x, rh_y, rh_vis = landmarks["right_hip"]
+            return ((lh_x + rh_x) / 2, (lh_y + rh_y) / 2, (lh_vis + rh_vis) / 2)
+        return (0.5, 0.5, 0.0)
+
+    # Normalize weights and compute weighted average
+    total_weight = sum(weights)
+    normalized_weights = [w / total_weight for w in weights]
+
+    com_x = float(sum(p[0] * w for p, w in zip(segments, normalized_weights, strict=True)))
+    com_y = float(sum(p[1] * w for p, w in zip(segments, normalized_weights, strict=True)))
+    com_visibility = float(np.mean(visibilities)) if visibilities else 0.0
+
+    return (com_x, com_y, com_visibility)
+
+
+def _add_head_segment(
+    segments: list,
+    weights: list,
+    visibilities: list,
+    landmarks: dict[str, tuple[float, float, float]],
+    vis_threshold: float,
+) -> None:
+    """Add head segment (8% body mass) if visible."""
     if "nose" in landmarks:
         x, y, vis = landmarks["nose"]
-        if vis > visibility_threshold:
+        if vis > vis_threshold:
             segments.append((x, y))
-            segment_weights.append(0.08)
+            weights.append(0.08)
             visibilities.append(vis)
 
-    # Trunk segment: 50% (midpoint between shoulders and hips)
-    trunk_landmarks = ["left_shoulder", "right_shoulder", "left_hip", "right_hip"]
-    trunk_positions = [
+
+def _add_trunk_segment(
+    segments: list,
+    weights: list,
+    visibilities: list,
+    landmarks: dict[str, tuple[float, float, float]],
+    vis_threshold: float,
+) -> None:
+    """Add trunk segment (50% body mass) if visible."""
+    trunk_keys = ["left_shoulder", "right_shoulder", "left_hip", "right_hip"]
+    trunk_pos = [
         (x, y, vis)
-        for key in trunk_landmarks
+        for key in trunk_keys
         if key in landmarks
         for x, y, vis in [landmarks[key]]
-        if vis > visibility_threshold
+        if vis > vis_threshold
     ]
-    if len(trunk_positions) >= 2:
-        trunk_x = float(np.mean([pos[0] for pos in trunk_positions]))
-        trunk_y = float(np.mean([pos[1] for pos in trunk_positions]))
-        trunk_vis = float(np.mean([pos[2] for pos in trunk_positions]))
+    if len(trunk_pos) >= 2:
+        trunk_x = float(np.mean([p[0] for p in trunk_pos]))
+        trunk_y = float(np.mean([p[1] for p in trunk_pos]))
+        trunk_vis = float(np.mean([p[2] for p in trunk_pos]))
         segments.append((trunk_x, trunk_y))
-        segment_weights.append(0.50)
+        weights.append(0.50)
         visibilities.append(trunk_vis)
 
-    # Thigh segment: 20% total (midpoint hip to knee for each leg)
-    for side in ["left", "right"]:
-        hip_key = f"{side}_hip"
-        knee_key = f"{side}_knee"
-        if hip_key in landmarks and knee_key in landmarks:
-            hip_x, hip_y, hip_vis = landmarks[hip_key]
-            knee_x, knee_y, knee_vis = landmarks[knee_key]
-            if hip_vis > visibility_threshold and knee_vis > visibility_threshold:
-                thigh_x = (hip_x + knee_x) / 2
-                thigh_y = (hip_y + knee_y) / 2
-                thigh_vis = (hip_vis + knee_vis) / 2
-                segments.append((thigh_x, thigh_y))
-                segment_weights.append(0.10)  # 10% per leg
-                visibilities.append(thigh_vis)
-
-    # Lower leg segment: 10% total (midpoint knee to ankle for each leg)
-    for side in ["left", "right"]:
-        knee_key = f"{side}_knee"
-        ankle_key = f"{side}_ankle"
-        if knee_key in landmarks and ankle_key in landmarks:
-            knee_x, knee_y, knee_vis = landmarks[knee_key]
-            ankle_x, ankle_y, ankle_vis = landmarks[ankle_key]
-            if knee_vis > visibility_threshold and ankle_vis > visibility_threshold:
-                leg_x = (knee_x + ankle_x) / 2
-                leg_y = (knee_y + ankle_y) / 2
-                leg_vis = (knee_vis + ankle_vis) / 2
-                segments.append((leg_x, leg_y))
-                segment_weights.append(0.05)  # 5% per leg
-                visibilities.append(leg_vis)
-
-    # Foot segment: 3% total (average of ankle, heel, foot_index)
-    for side in ["left", "right"]:
-        foot_keys = [f"{side}_ankle", f"{side}_heel", f"{side}_foot_index"]
-        foot_positions = [
-            (x, y, vis)
-            for key in foot_keys
-            if key in landmarks
-            for x, y, vis in [landmarks[key]]
-            if vis > visibility_threshold
-        ]
-        if foot_positions:
-            foot_x = float(np.mean([pos[0] for pos in foot_positions]))
-            foot_y = float(np.mean([pos[1] for pos in foot_positions]))
-            foot_vis = float(np.mean([pos[2] for pos in foot_positions]))
-            segments.append((foot_x, foot_y))
-            segment_weights.append(0.015)  # 1.5% per foot
-            visibilities.append(foot_vis)
-
-    # If no segments found, fall back to hip average
-    if not segments:
-        if "left_hip" in landmarks and "right_hip" in landmarks:
-            lh_x, lh_y, lh_vis = landmarks["left_hip"]
-            rh_x, rh_y, rh_vis = landmarks["right_hip"]
-            return (
-                (lh_x + rh_x) / 2,
-                (lh_y + rh_y) / 2,
-                (lh_vis + rh_vis) / 2,
-            )
-        # Ultimate fallback: center of frame
-        return (0.5, 0.5, 0.0)
 
-    # Normalize weights to sum to 1.0
-    total_weight = sum(segment_weights)
-    normalized_weights = [w / total_weight for w in segment_weights]
+def _add_limb_segment(
+    segments: list,
+    weights: list,
+    visibilities: list,
+    landmarks: dict[str, tuple[float, float, float]],
+    side: str,
+    proximal_key: str,
+    distal_key: str,
+    segment_weight: float,
+    vis_threshold: float,
+) -> None:
+    """Add a limb segment (thigh or lower leg) if both endpoints visible."""
+    prox_full = f"{side}_{proximal_key}"
+    dist_full = f"{side}_{distal_key}"
 
-    # Compute weighted average of segment positions
-    com_x = float(
-        sum(
-            pos[0] * weight
-            for pos, weight in zip(segments, normalized_weights, strict=True)
-        )
-    )
-    com_y = float(
-        sum(
-            pos[1] * weight
-            for pos, weight in zip(segments, normalized_weights, strict=True)
-        )
-    )
-    com_visibility = float(np.mean(visibilities)) if visibilities else 0.0
+    if prox_full in landmarks and dist_full in landmarks:
+        px, py, pvis = landmarks[prox_full]
+        dx, dy, dvis = landmarks[dist_full]
+        if pvis > vis_threshold and dvis > vis_threshold:
+            seg_x = (px + dx) / 2
+            seg_y = (py + dy) / 2
+            seg_vis = (pvis + dvis) / 2
+            segments.append((seg_x, seg_y))
+            weights.append(segment_weight)
+            visibilities.append(seg_vis)
 
-    return (com_x, com_y, com_visibility)
+
+def _add_foot_segment(
+    segments: list,
+    weights: list,
+    visibilities: list,
+    landmarks: dict[str, tuple[float, float, float]],
+    side: str,
+    vis_threshold: float,
+) -> None:
+    """Add foot segment (1.5% body mass per foot) if visible."""
+    foot_keys = [f"{side}_ankle", f"{side}_heel", f"{side}_foot_index"]
+    foot_pos = [
+        (x, y, vis)
+        for key in foot_keys
+        if key in landmarks
+        for x, y, vis in [landmarks[key]]
+        if vis > vis_threshold
+    ]
+    if foot_pos:
+        foot_x = float(np.mean([p[0] for p in foot_pos]))
+        foot_y = float(np.mean([p[1] for p in foot_pos]))
+        foot_vis = float(np.mean([p[2] for p in foot_pos]))
+        segments.append((foot_x, foot_y))
+        weights.append(0.015)
+        visibilities.append(foot_vis)

kinemotion/core/pose_landmarks.py

@@ -0,0 +1,67 @@
+"""MediaPipe Pose landmark index constants for Tasks API.
+
+The MediaPipe Tasks API uses index-based landmark access (0-32) instead of enums.
+This module provides named constants for the 33 pose landmarks.
+"""
+
+from __future__ import annotations
+
+# MediaPipe Pose has 33 landmarks per pose
+# See: https://ai.google.dev/edge/mediapipe/solutions/vision/pose_landmarker
+LANDMARK_INDICES: dict[str, int] = {
+    # Face
+    "nose": 0,
+    "left_eye_inner": 1,
+    "left_eye": 2,
+    "left_eye_outer": 3,
+    "right_eye_inner": 4,
+    "right_eye": 5,
+    "right_eye_outer": 6,
+    "left_ear": 7,
+    "right_ear": 8,
+    "mouth_left": 9,
+    "mouth_right": 10,
+    # Upper body
+    "left_shoulder": 11,
+    "right_shoulder": 12,
+    "left_elbow": 13,
+    "right_elbow": 14,
+    "left_wrist": 15,
+    "right_wrist": 16,
+    "left_pinky": 17,
+    "right_pinky": 18,
+    "left_index": 19,
+    "right_index": 20,
+    "left_thumb": 21,
+    "right_thumb": 22,
+    # Lower body
+    "left_hip": 23,
+    "right_hip": 24,
+    "left_knee": 25,
+    "right_knee": 26,
+    "left_ankle": 27,
+    "right_ankle": 28,
+    "left_heel": 29,
+    "right_heel": 30,
+    "left_foot_index": 31,
+    "right_foot_index": 32,
+}
+
+# Landmarks used in kinemotion analysis
+KINEMOTION_LANDMARKS = {
+    "nose",
+    "left_shoulder",
+    "right_shoulder",
+    "left_hip",
+    "right_hip",
+    "left_knee",
+    "right_knee",
+    "left_ankle",
+    "right_ankle",
+    "left_heel",
+    "right_heel",
+    "left_foot_index",
+    "right_foot_index",
+}
+
+__all__ = ["LANDMARK_INDICES", "KINEMOTION_LANDMARKS"]