kinemotion 0.17.0__py3-none-any.whl

kinemotion/cmj/debug_overlay.py

@@ -0,0 +1,463 @@
+"""Debug overlay rendering for CMJ analysis."""
+
+import cv2
+import numpy as np
+
+from ..core.debug_overlay_utils import BaseDebugOverlayRenderer
+from .joint_angles import calculate_triple_extension
+from .kinematics import CMJMetrics
+
+
+class CMJPhaseState:
+    """States for CMJ phases."""
+
+    STANDING = "standing"
+    ECCENTRIC = "eccentric"
+    TRANSITION = "transition"
+    CONCENTRIC = "concentric"
+    FLIGHT = "flight"
+    LANDING = "landing"
+
+
+class CMJDebugOverlayRenderer(BaseDebugOverlayRenderer):
+    """Renders debug information on CMJ video frames."""
+
+    def _determine_phase(self, frame_idx: int, metrics: CMJMetrics) -> str:
+        """Determine which phase the current frame is in."""
+        if metrics.standing_start_frame and frame_idx < metrics.standing_start_frame:
+            return CMJPhaseState.STANDING
+
+        if frame_idx < metrics.lowest_point_frame:
+            return CMJPhaseState.ECCENTRIC
+
+        # Brief transition at lowest point (±2 frames)
+        if abs(frame_idx - metrics.lowest_point_frame) < 2:
+            return CMJPhaseState.TRANSITION
+
+        if frame_idx < metrics.takeoff_frame:
+            return CMJPhaseState.CONCENTRIC
+
+        if frame_idx < metrics.landing_frame:
+            return CMJPhaseState.FLIGHT
+
+        return CMJPhaseState.LANDING
+
+    def _get_phase_color(self, phase: str) -> tuple[int, int, int]:
+        """Get color for each phase."""
+        colors = {
+            CMJPhaseState.STANDING: (255, 200, 100),  # Light blue
+            CMJPhaseState.ECCENTRIC: (0, 165, 255),  # Orange
+            CMJPhaseState.TRANSITION: (255, 0, 255),  # Magenta/Purple
+            CMJPhaseState.CONCENTRIC: (0, 255, 0),  # Green
+            CMJPhaseState.FLIGHT: (0, 0, 255),  # Red
+            CMJPhaseState.LANDING: (255, 255, 255),  # White
+        }
+        return colors.get(phase, (128, 128, 128))
+
+    def _get_skeleton_segments(
+        self, side_prefix: str
+    ) -> list[tuple[str, str, tuple[int, int, int], int]]:
+        """Get skeleton segments for one side of the body."""
+        return [
+            (f"{side_prefix}heel", f"{side_prefix}ankle", (0, 255, 255), 3),  # Foot
+            (
+                f"{side_prefix}heel",
+                f"{side_prefix}foot_index",
+                (0, 255, 255),
+                2,
+            ),  # Alt foot
+            (f"{side_prefix}ankle", f"{side_prefix}knee", (255, 100, 100), 4),  # Shin
+            (f"{side_prefix}knee", f"{side_prefix}hip", (100, 255, 100), 4),  # Femur
+            (
+                f"{side_prefix}hip",
+                f"{side_prefix}shoulder",
+                (100, 100, 255),
+                4,
+            ),  # Trunk
+            (f"{side_prefix}shoulder", "nose", (150, 150, 255), 2),  # Neck
+        ]
+
+    def _draw_segment(
+        self,
+        frame: np.ndarray,
+        landmarks: dict[str, tuple[float, float, float]],
+        start_key: str,
+        end_key: str,
+        color: tuple[int, int, int],
+        thickness: int,
+    ) -> None:
+        """Draw a single skeleton segment if both endpoints are visible."""
+        if start_key not in landmarks or end_key not in landmarks:
+            return
+
+        start_vis = landmarks[start_key][2]
+        end_vis = landmarks[end_key][2]
+
+        # Very low threshold to show as much as possible
+        if start_vis > 0.2 and end_vis > 0.2:
+            start_x = int(landmarks[start_key][0] * self.width)
+            start_y = int(landmarks[start_key][1] * self.height)
+            end_x = int(landmarks[end_key][0] * self.width)
+            end_y = int(landmarks[end_key][1] * self.height)
+
+            cv2.line(frame, (start_x, start_y), (end_x, end_y), color, thickness)
+
+    def _draw_joints(
+        self,
+        frame: np.ndarray,
+        landmarks: dict[str, tuple[float, float, float]],
+        side_prefix: str,
+    ) -> None:
+        """Draw joint circles for one side of the body."""
+        joint_keys = [
+            f"{side_prefix}heel",
+            f"{side_prefix}foot_index",
+            f"{side_prefix}ankle",
+            f"{side_prefix}knee",
+            f"{side_prefix}hip",
+            f"{side_prefix}shoulder",
+        ]
+        for key in joint_keys:
+            if key in landmarks and landmarks[key][2] > 0.2:
+                jx = int(landmarks[key][0] * self.width)
+                jy = int(landmarks[key][1] * self.height)
+                cv2.circle(frame, (jx, jy), 6, (255, 255, 255), -1)
+                cv2.circle(frame, (jx, jy), 8, (0, 0, 0), 2)
+
+    def _draw_skeleton(
+        self, frame: np.ndarray, landmarks: dict[str, tuple[float, float, float]]
+    ) -> None:
+        """Draw skeleton segments showing body landmarks.
+
+        Draws whatever landmarks are visible. In side-view videos, ankle/knee
+        may have low visibility, so we draw available segments.
+
+        Args:
+            frame: Frame to draw on (modified in place)
+            landmarks: Pose landmarks
+        """
+        # Try both sides and draw all visible segments
+        for side_prefix in ["right_", "left_"]:
+            segments = self._get_skeleton_segments(side_prefix)
+
+            # Draw ALL visible segments (not just one side)
+            for start_key, end_key, color, thickness in segments:
+                self._draw_segment(
+                    frame, landmarks, start_key, end_key, color, thickness
+                )
+
+            # Draw joints as circles for this side
+            self._draw_joints(frame, landmarks, side_prefix)
+
+        # Always draw nose (head position) if visible
+        if "nose" in landmarks and landmarks["nose"][2] > 0.2:
+            nx = int(landmarks["nose"][0] * self.width)
+            ny = int(landmarks["nose"][1] * self.height)
+            cv2.circle(frame, (nx, ny), 8, (255, 255, 0), -1)
+            cv2.circle(frame, (nx, ny), 10, (0, 0, 0), 2)
+
+    def _draw_joint_angles(
+        self,
+        frame: np.ndarray,
+        landmarks: dict[str, tuple[float, float, float]],
+        phase_color: tuple[int, int, int],
+    ) -> None:
+        """Draw joint angles for triple extension analysis.
+
+        Args:
+            frame: Frame to draw on (modified in place)
+            landmarks: Pose landmarks
+            phase_color: Current phase color
+        """
+        # Try right side first, fallback to left
+        angles = calculate_triple_extension(landmarks, side="right")
+        side_used = "right"
+
+        if angles is None:
+            angles = calculate_triple_extension(landmarks, side="left")
+            side_used = "left"
+
+        if angles is None:
+            return
+
+        # Position for angle text display (right side of frame)
+        text_x = self.width - 180
+        text_y = 100
+
+        # Draw background box for angles
+        box_height = 150
+        cv2.rectangle(
+            frame,
+            (text_x - 10, text_y - 30),
+            (self.width - 10, text_y + box_height),
+            (0, 0, 0),
+            -1,
+        )
+        cv2.rectangle(
+            frame,
+            (text_x - 10, text_y - 30),
+            (self.width - 10, text_y + box_height),
+            phase_color,
+            2,
+        )
+
+        # Title
+        cv2.putText(
+            frame,
+            "TRIPLE EXTENSION",
+            (text_x, text_y - 5),
+            cv2.FONT_HERSHEY_SIMPLEX,
+            0.5,
+            (255, 255, 255),
+            1,
+        )
+
+        # Draw available angles (show "N/A" for unavailable)
+        angle_data = [
+            ("Ankle", angles.get("ankle_angle"), (0, 255, 255)),
+            ("Knee", angles.get("knee_angle"), (255, 100, 100)),
+            ("Hip", angles.get("hip_angle"), (100, 255, 100)),
+            ("Trunk", angles.get("trunk_tilt"), (100, 100, 255)),
+        ]
+
+        y_offset = text_y + 25
+        for label, angle, color in angle_data:
+            # Angle text
+            if angle is not None:
+                angle_text = f"{label}: {angle:.0f}"
+                text_color = color
+            else:
+                angle_text = f"{label}: N/A"
+                text_color = (128, 128, 128)  # Gray for unavailable
+
+            cv2.putText(
+                frame,
+                angle_text,
+                (text_x, y_offset),
+                cv2.FONT_HERSHEY_SIMPLEX,
+                0.5,
+                text_color,
+                2,
+            )
+            y_offset += 30
+
+        # Draw angle arcs at joints for visual feedback (only if angle is available)
+        ankle_angle = angles.get("ankle_angle")
+        if ankle_angle is not None:
+            self._draw_angle_arc(frame, landmarks, f"{side_used}_ankle", ankle_angle)
+        knee_angle = angles.get("knee_angle")
+        if knee_angle is not None:
+            self._draw_angle_arc(frame, landmarks, f"{side_used}_knee", knee_angle)
+        hip_angle = angles.get("hip_angle")
+        if hip_angle is not None:
+            self._draw_angle_arc(frame, landmarks, f"{side_used}_hip", hip_angle)
+
+    def _draw_angle_arc(
+        self,
+        frame: np.ndarray,
+        landmarks: dict[str, tuple[float, float, float]],
+        joint_key: str,
+        angle: float,
+    ) -> None:
+        """Draw a small arc at a joint to visualize the angle.
+
+        Args:
+            frame: Frame to draw on (modified in place)
+            landmarks: Pose landmarks
+            joint_key: Key of the joint landmark
+            angle: Angle value in degrees
+        """
+        if joint_key not in landmarks or landmarks[joint_key][2] < 0.3:
+            return
+
+        jx = int(landmarks[joint_key][0] * self.width)
+        jy = int(landmarks[joint_key][1] * self.height)
+
+        # Draw arc radius based on angle (smaller arc for more extended joints)
+        radius = 25
+
+        # Color based on extension: green for extended (>160°), red for flexed (<90°)
+        if angle > 160:
+            arc_color = (0, 255, 0)  # Green - good extension
+        elif angle < 90:
+            arc_color = (0, 0, 255)  # Red - deep flexion
+        else:
+            arc_color = (0, 165, 255)  # Orange - moderate
+
+        # Draw arc (simplified as a circle for now)
+        cv2.circle(frame, (jx, jy), radius, arc_color, 2)
+
+    def _draw_foot_landmarks(
+        self,
+        frame: np.ndarray,
+        landmarks: dict[str, tuple[float, float, float]],
+        phase_color: tuple[int, int, int],
+    ) -> None:
+        """Draw foot landmarks and average position."""
+        foot_keys = ["left_ankle", "right_ankle", "left_heel", "right_heel"]
+        foot_positions = []
+
+        for key in foot_keys:
+            if key in landmarks:
+                x, y, vis = landmarks[key]
+                if vis > 0.5:
+                    lx = int(x * self.width)
+                    ly = int(y * self.height)
+                    foot_positions.append((lx, ly))
+                    cv2.circle(frame, (lx, ly), 5, (255, 255, 0), -1)
+
+        # Draw average foot position with phase color
+        if foot_positions:
+            avg_x = int(np.mean([p[0] for p in foot_positions]))
+            avg_y = int(np.mean([p[1] for p in foot_positions]))
+            cv2.circle(frame, (avg_x, avg_y), 12, phase_color, -1)
+            cv2.circle(frame, (avg_x, avg_y), 14, (255, 255, 255), 2)
+
+    def _draw_phase_banner(
+        self, frame: np.ndarray, phase: str | None, phase_color: tuple[int, int, int]
+    ) -> None:
+        """Draw phase indicator banner."""
+        if not phase:
+            return
+
+        phase_text = f"Phase: {phase.upper()}"
+        text_size = cv2.getTextSize(phase_text, cv2.FONT_HERSHEY_SIMPLEX, 1, 2)[0]
+        cv2.rectangle(frame, (5, 5), (text_size[0] + 15, 45), phase_color, -1)
+        cv2.putText(
+            frame, phase_text, (10, 35), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 0), 2
+        )
+
+    def _draw_key_frame_markers(
+        self, frame: np.ndarray, frame_idx: int, metrics: CMJMetrics
+    ) -> None:
+        """Draw markers for key frames (standing start, lowest, takeoff, landing)."""
+        y_offset = 120
+        markers = []
+
+        if metrics.standing_start_frame and frame_idx == int(
+            metrics.standing_start_frame
+        ):
+            markers.append("COUNTERMOVEMENT START")
+
+        if frame_idx == int(metrics.lowest_point_frame):
+            markers.append("LOWEST POINT")
+
+        if frame_idx == int(metrics.takeoff_frame):
+            markers.append("TAKEOFF")
+
+        if frame_idx == int(metrics.landing_frame):
+            markers.append("LANDING")
+
+        for marker in markers:
+            cv2.putText(
+                frame,
+                marker,
+                (10, y_offset),
+                cv2.FONT_HERSHEY_SIMPLEX,
+                0.7,
+                (255, 255, 0),
+                2,
+            )
+            y_offset += 35
+
+    def _draw_metrics_summary(
+        self, frame: np.ndarray, frame_idx: int, metrics: CMJMetrics
+    ) -> None:
+        """Draw metrics summary in bottom right (last 30 frames)."""
+        total_frames = int(metrics.landing_frame) + 30
+        if frame_idx < total_frames - 30:
+            return
+
+        metrics_text = [
+            f"Jump Height: {metrics.jump_height:.3f}m",
+            f"Flight Time: {metrics.flight_time*1000:.0f}ms",
+            f"CM Depth: {metrics.countermovement_depth:.3f}m",
+            f"Ecc Duration: {metrics.eccentric_duration*1000:.0f}ms",
+            f"Con Duration: {metrics.concentric_duration*1000:.0f}ms",
+        ]
+
+        # Draw background
+        box_height = len(metrics_text) * 30 + 20
+        cv2.rectangle(
+            frame,
+            (self.width - 320, self.height - box_height - 10),
+            (self.width - 10, self.height - 10),
+            (0, 0, 0),
+            -1,
+        )
+        cv2.rectangle(
+            frame,
+            (self.width - 320, self.height - box_height - 10),
+            (self.width - 10, self.height - 10),
+            (0, 255, 0),
+            2,
+        )
+
+        # Draw metrics text
+        text_y = self.height - box_height + 10
+        for text in metrics_text:
+            cv2.putText(
+                frame,
+                text,
+                (self.width - 310, text_y),
+                cv2.FONT_HERSHEY_SIMPLEX,
+                0.6,
+                (255, 255, 255),
+                1,
+            )
+            text_y += 30
+
+    def render_frame(
+        self,
+        frame: np.ndarray,
+        landmarks: dict[str, tuple[float, float, float]] | None,
+        frame_idx: int,
+        metrics: CMJMetrics | None = None,
+    ) -> np.ndarray:
+        """
+        Render debug overlay on frame.
+
+        Args:
+            frame: Original video frame
+            landmarks: Pose landmarks for this frame
+            frame_idx: Current frame index
+            metrics: CMJ metrics (optional)
+
+        Returns:
+            Frame with debug overlay
+        """
+        annotated = frame.copy()
+
+        # Determine current phase if metrics available
+        phase = None
+        phase_color = (255, 255, 255)
+        if metrics:
+            phase = self._determine_phase(frame_idx, metrics)
+            phase_color = self._get_phase_color(phase)
+
+        # Draw skeleton and triple extension if landmarks available
+        if landmarks:
+            self._draw_skeleton(annotated, landmarks)
+            self._draw_joint_angles(annotated, landmarks, phase_color)
+            self._draw_foot_landmarks(annotated, landmarks, phase_color)
+
+        # Draw phase indicator banner
+        self._draw_phase_banner(annotated, phase, phase_color)
+
+        # Draw frame number
+        cv2.putText(
+            annotated,
+            f"Frame: {frame_idx}",
+            (10, 80),
+            cv2.FONT_HERSHEY_SIMPLEX,
+            0.7,
+            (255, 255, 255),
+            2,
+        )
+
+        # Draw key frame markers and metrics summary
+        if metrics:
+            self._draw_key_frame_markers(annotated, frame_idx, metrics)
+            self._draw_metrics_summary(annotated, frame_idx, metrics)
+
+        return annotated