kinemotion 0.10.2__py3-none-any.whl

kinemotion/core/pose.py

@@ -0,0 +1,220 @@
+"""Pose tracking using MediaPipe Pose."""
+
+import cv2
+import mediapipe as mp
+import numpy as np
+
+
+class PoseTracker:
+    """Tracks human pose landmarks in video frames using MediaPipe."""
+
+    def __init__(
+        self,
+        min_detection_confidence: float = 0.5,
+        min_tracking_confidence: float = 0.5,
+    ):
+        """
+        Initialize the pose tracker.
+
+        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,
+            min_tracking_confidence=min_tracking_confidence,
+            model_complexity=1,
+        )
+
+    def process_frame(
+        self, frame: np.ndarray
+    ) -> dict[str, tuple[float, float, float]] | None:
+        """
+        Process a single frame and extract pose landmarks.
+
+        Args:
+            frame: BGR image frame
+
+        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)
+
+        # Process the frame
+        results = self.pose.process(rgb_frame)
+
+        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)
+
+        return landmarks
+
+    def close(self) -> None:
+        """Release resources."""
+        self.pose.close()
+
+
+def compute_center_of_mass(
+    landmarks: dict[str, tuple[float, float, float]],
+    visibility_threshold: float = 0.5,
+) -> tuple[float, float, float]:
+    """
+    Compute approximate center of mass (CoM) from body landmarks.
+
+    Uses biomechanical segment weights based on Dempster's body segment parameters:
+    - Head: 8% of body mass (represented by nose)
+    - Trunk (shoulders to hips): 50% of body mass
+    - Thighs: 2 × 10% = 20% of body mass
+    - Legs (knees to ankles): 2 × 5% = 10% of body mass
+    - Feet: 2 × 1.5% = 3% of body mass
+
+    The CoM is estimated as a weighted average of these segments, with
+    weights corresponding to their proportion of total body mass.
+
+    Args:
+        landmarks: Dictionary of landmark positions (x, y, visibility)
+        visibility_threshold: Minimum visibility to include landmark in calculation
+
+    Returns:
+        (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 = []
+
+    # Head segment: 8% (use nose as proxy)
+    if "nose" in landmarks:
+        x, y, vis = landmarks["nose"]
+        if vis > visibility_threshold:
+            segments.append((x, y))
+            segment_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 = [
+        (x, y, vis)
+        for key in trunk_landmarks
+        if key in landmarks
+        for x, y, vis in [landmarks[key]]
+        if vis > visibility_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]))
+        segments.append((trunk_x, trunk_y))
+        segment_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]
+
+    # 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
+
+    return (com_x, com_y, com_visibility)

kinemotion/core/smoothing.py

@@ -0,0 +1,366 @@
+"""Landmark smoothing utilities to reduce jitter in pose tracking."""
+
+import numpy as np
+from scipy.signal import savgol_filter
+
+from .filtering import (
+    bilateral_temporal_filter,
+    reject_outliers,
+)
+
+
+def smooth_landmarks(
+    landmark_sequence: list[dict[str, tuple[float, float, float]] | None],
+    window_length: int = 5,
+    polyorder: int = 2,
+) -> list[dict[str, tuple[float, float, float]] | None]:
+    """
+    Smooth landmark trajectories using Savitzky-Golay filter.
+
+    Args:
+        landmark_sequence: List of landmark dictionaries from each frame
+        window_length: Length of filter window (must be odd, >= polyorder + 2)
+        polyorder: Order of polynomial used to fit samples
+
+    Returns:
+        Smoothed landmark sequence with same structure as input
+    """
+    if len(landmark_sequence) < window_length:
+        # Not enough frames to smooth effectively
+        return landmark_sequence
+
+    # Ensure window_length is odd
+    if window_length % 2 == 0:
+        window_length += 1
+
+    # Extract landmark names from first valid frame
+    landmark_names = None
+    for frame_landmarks in landmark_sequence:
+        if frame_landmarks is not None:
+            landmark_names = list(frame_landmarks.keys())
+            break
+
+    if landmark_names is None:
+        return landmark_sequence
+
+    # Build arrays for each landmark coordinate
+    smoothed_sequence: list[dict[str, tuple[float, float, float]] | None] = []
+
+    for landmark_name in landmark_names:
+        # Extract x, y coordinates for this landmark across all frames
+        x_coords = []
+        y_coords = []
+        valid_frames = []
+
+        for i, frame_landmarks in enumerate(landmark_sequence):
+            if frame_landmarks is not None and landmark_name in frame_landmarks:
+                x, y, _ = frame_landmarks[landmark_name]  # vis not used
+                x_coords.append(x)
+                y_coords.append(y)
+                valid_frames.append(i)
+
+        if len(x_coords) < window_length:
+            continue
+
+        # Apply Savitzky-Golay filter
+        x_smooth = savgol_filter(x_coords, window_length, polyorder)
+        y_smooth = savgol_filter(y_coords, window_length, polyorder)
+
+        # Store smoothed values back
+        for idx, frame_idx in enumerate(valid_frames):
+            if frame_idx >= len(smoothed_sequence):
+                smoothed_sequence.extend(
+                    [{}] * (frame_idx - len(smoothed_sequence) + 1)
+                )
+
+            # Ensure smoothed_sequence[frame_idx] is a dict, not None
+            if smoothed_sequence[frame_idx] is None:
+                smoothed_sequence[frame_idx] = {}
+
+            if (
+                landmark_name not in smoothed_sequence[frame_idx]
+                and landmark_sequence[frame_idx] is not None
+            ):
+                # Keep original visibility
+                orig_vis = landmark_sequence[frame_idx][landmark_name][2]
+                smoothed_sequence[frame_idx][landmark_name] = (
+                    float(x_smooth[idx]),
+                    float(y_smooth[idx]),
+                    orig_vis,
+                )
+
+    # Fill in any missing frames with original data
+    for i in range(len(landmark_sequence)):
+        if i >= len(smoothed_sequence) or not smoothed_sequence[i]:
+            if i < len(smoothed_sequence):
+                smoothed_sequence[i] = landmark_sequence[i]
+            else:
+                smoothed_sequence.append(landmark_sequence[i])
+
+    return smoothed_sequence
+
+
+def compute_velocity(
+    positions: np.ndarray, fps: float, smooth_window: int = 3
+) -> np.ndarray:
+    """
+    Compute velocity from position data.
+
+    Args:
+        positions: Array of positions over time (n_frames, n_dims)
+        fps: Frames per second of the video
+        smooth_window: Window size for velocity smoothing
+
+    Returns:
+        Velocity array (n_frames, n_dims)
+    """
+    dt = 1.0 / fps
+    velocity = np.gradient(positions, dt, axis=0)
+
+    # Smooth velocity if we have enough data
+    if len(velocity) >= smooth_window and smooth_window > 1:
+        if smooth_window % 2 == 0:
+            smooth_window += 1
+        for dim in range(velocity.shape[1]):
+            velocity[:, dim] = savgol_filter(velocity[:, dim], smooth_window, 1)
+
+    return velocity
+
+
+def compute_velocity_from_derivative(
+    positions: np.ndarray,
+    window_length: int = 5,
+    polyorder: int = 2,
+) -> np.ndarray:
+    """
+    Compute velocity as derivative of smoothed position trajectory.
+
+    Uses Savitzky-Golay filter to compute the derivative directly, which provides
+    a much smoother and more accurate velocity estimate than frame-to-frame differences.
+
+    This method:
+    1. Fits a polynomial to the position data in a sliding window
+    2. Analytically computes the derivative of that polynomial
+    3. Returns smooth velocity values
+
+    Args:
+        positions: 1D array of position values (e.g., foot y-positions)
+        window_length: Window size for smoothing (must be odd, >= polyorder + 2)
+        polyorder: Polynomial order for Savitzky-Golay filter (typically 2 or 3)
+
+    Returns:
+        Array of absolute velocity values (magnitude of derivative)
+    """
+    if len(positions) < window_length:
+        # Fallback to simple differences for short sequences
+        return np.abs(np.diff(positions, prepend=positions[0]))
+
+    # Ensure window_length is odd
+    if window_length % 2 == 0:
+        window_length += 1
+
+    # Compute derivative using Savitzky-Golay filter
+    # deriv=1: compute first derivative
+    # delta=1.0: frame spacing (velocity per frame)
+    # mode='interp': interpolate at boundaries
+    velocity = savgol_filter(
+        positions,
+        window_length,
+        polyorder,
+        deriv=1,  # First derivative
+        delta=1.0,  # Frame spacing
+        mode="interp",
+    )
+
+    # Return absolute velocity (magnitude only)
+    return np.abs(velocity)
+
+
+def compute_acceleration_from_derivative(
+    positions: np.ndarray,
+    window_length: int = 5,
+    polyorder: int = 2,
+) -> np.ndarray:
+    """
+    Compute acceleration as second derivative of smoothed position trajectory.
+
+    Uses Savitzky-Golay filter to compute the second derivative directly,
+    providing smooth acceleration (curvature) estimates for detecting
+    characteristic patterns at landing and takeoff.
+
+    Landing and takeoff events show distinctive acceleration patterns:
+    - Landing: Large acceleration spike as feet decelerate on impact
+    - Takeoff: Acceleration change as body accelerates upward
+    - In flight: Constant acceleration due to gravity
+    - On ground: Near-zero acceleration (stationary position)
+
+    Args:
+        positions: 1D array of position values (e.g., foot y-positions)
+        window_length: Window size for smoothing (must be odd, >= polyorder + 2)
+        polyorder: Polynomial order for Savitzky-Golay filter (typically 2 or 3)
+
+    Returns:
+        Array of acceleration values (second derivative of position)
+    """
+    if len(positions) < window_length:
+        # Fallback to simple second differences for short sequences
+        velocity = np.diff(positions, prepend=positions[0])
+        return np.diff(velocity, prepend=velocity[0])
+
+    # Ensure window_length is odd
+    if window_length % 2 == 0:
+        window_length += 1
+
+    # Compute second derivative using Savitzky-Golay filter
+    # deriv=2: compute second derivative (acceleration/curvature)
+    # delta=1.0: frame spacing
+    # mode='interp': interpolate at boundaries
+    acceleration = savgol_filter(
+        positions,
+        window_length,
+        polyorder,
+        deriv=2,  # Second derivative
+        delta=1.0,  # Frame spacing
+        mode="interp",
+    )
+
+    return acceleration
+
+
+def smooth_landmarks_advanced(
+    landmark_sequence: list[dict[str, tuple[float, float, float]] | None],
+    window_length: int = 5,
+    polyorder: int = 2,
+    use_outlier_rejection: bool = True,
+    use_bilateral: bool = False,
+    ransac_threshold: float = 0.02,
+    bilateral_sigma_spatial: float = 3.0,
+    bilateral_sigma_intensity: float = 0.02,
+) -> list[dict[str, tuple[float, float, float]] | None]:
+    """
+    Advanced landmark smoothing with outlier rejection and bilateral filtering.
+
+    Combines multiple techniques for robust smoothing:
+    1. Outlier rejection (RANSAC + median filtering)
+    2. Optional bilateral filtering (edge-preserving)
+    3. Savitzky-Golay smoothing
+
+    Args:
+        landmark_sequence: List of landmark dictionaries from each frame
+        window_length: Length of filter window (must be odd, >= polyorder + 2)
+        polyorder: Order of polynomial used to fit samples
+        use_outlier_rejection: Apply outlier detection and removal
+        use_bilateral: Use bilateral filter instead of Savitzky-Golay
+        ransac_threshold: Threshold for RANSAC outlier detection
+        bilateral_sigma_spatial: Spatial sigma for bilateral filter
+        bilateral_sigma_intensity: Intensity sigma for bilateral filter
+
+    Returns:
+        Smoothed landmark sequence with same structure as input
+    """
+    if len(landmark_sequence) < window_length:
+        # Not enough frames to smooth effectively
+        return landmark_sequence
+
+    # Ensure window_length is odd
+    if window_length % 2 == 0:
+        window_length += 1
+
+    # Extract landmark names from first valid frame
+    landmark_names = None
+    for frame_landmarks in landmark_sequence:
+        if frame_landmarks is not None:
+            landmark_names = list(frame_landmarks.keys())
+            break
+
+    if landmark_names is None:
+        return landmark_sequence
+
+    # Build arrays for each landmark coordinate
+    smoothed_sequence: list[dict[str, tuple[float, float, float]] | None] = []
+
+    for landmark_name in landmark_names:
+        # Extract x, y coordinates for this landmark across all frames
+        x_coords = []
+        y_coords = []
+        valid_frames = []
+
+        for i, frame_landmarks in enumerate(landmark_sequence):
+            if frame_landmarks is not None and landmark_name in frame_landmarks:
+                x, y, _ = frame_landmarks[landmark_name]  # vis not used
+                x_coords.append(x)
+                y_coords.append(y)
+                valid_frames.append(i)
+
+        if len(x_coords) < window_length:
+            continue
+
+        x_array = np.array(x_coords)
+        y_array = np.array(y_coords)
+
+        # Step 1: Outlier rejection
+        if use_outlier_rejection:
+            x_array, _ = reject_outliers(
+                x_array,
+                use_ransac=True,
+                use_median=True,
+                ransac_threshold=ransac_threshold,
+            )
+            y_array, _ = reject_outliers(
+                y_array,
+                use_ransac=True,
+                use_median=True,
+                ransac_threshold=ransac_threshold,
+            )
+
+        # Step 2: Smoothing (bilateral or Savitzky-Golay)
+        if use_bilateral:
+            x_smooth = bilateral_temporal_filter(
+                x_array,
+                window_size=window_length,
+                sigma_spatial=bilateral_sigma_spatial,
+                sigma_intensity=bilateral_sigma_intensity,
+            )
+            y_smooth = bilateral_temporal_filter(
+                y_array,
+                window_size=window_length,
+                sigma_spatial=bilateral_sigma_spatial,
+                sigma_intensity=bilateral_sigma_intensity,
+            )
+        else:
+            # Standard Savitzky-Golay
+            x_smooth = savgol_filter(x_array, window_length, polyorder)
+            y_smooth = savgol_filter(y_array, window_length, polyorder)
+
+        # Store smoothed values back
+        for idx, frame_idx in enumerate(valid_frames):
+            if frame_idx >= len(smoothed_sequence):
+                smoothed_sequence.extend(
+                    [{}] * (frame_idx - len(smoothed_sequence) + 1)
+                )
+
+            # Ensure smoothed_sequence[frame_idx] is a dict, not None
+            if smoothed_sequence[frame_idx] is None:
+                smoothed_sequence[frame_idx] = {}
+
+            if (
+                landmark_name not in smoothed_sequence[frame_idx]
+                and landmark_sequence[frame_idx] is not None
+            ):
+                # Keep original visibility
+                orig_vis = landmark_sequence[frame_idx][landmark_name][2]
+                smoothed_sequence[frame_idx][landmark_name] = (
+                    float(x_smooth[idx]),
+                    float(y_smooth[idx]),
+                    orig_vis,
+                )
+
+    # Fill in any missing frames with original data
+    for i in range(len(landmark_sequence)):
+        if i >= len(smoothed_sequence) or not smoothed_sequence[i]:
+            if i < len(smoothed_sequence):
+                smoothed_sequence[i] = landmark_sequence[i]
+            else:
+                smoothed_sequence.append(landmark_sequence[i])
+
+    return smoothed_sequence