kinemotion 0.10.12__py3-none-any.whl → 0.11.0__py3-none-any.whl

kinemotion/cmj/debug_overlay.py

@@ -0,0 +1,514 @@
+"""Debug overlay rendering for CMJ analysis."""
+
+import cv2
+import numpy as np
+
+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:
+    """Renders debug information on CMJ video frames."""
+
+    def __init__(
+        self,
+        output_path: str,
+        width: int,
+        height: int,
+        display_width: int,
+        display_height: int,
+        fps: float,
+    ):
+        """
+        Initialize overlay renderer.
+
+        Args:
+            output_path: Path for output video
+            width: Encoded frame width (from source video)
+            height: Encoded frame height (from source video)
+            display_width: Display width (considering SAR)
+            display_height: Display height (considering SAR)
+            fps: Frames per second
+        """
+        self.width = width
+        self.height = height
+        self.display_width = display_width
+        self.display_height = display_height
+        self.needs_resize = (display_width != width) or (display_height != height)
+
+        # Try H.264 codec first (better quality/compatibility), fallback to mp4v
+        fourcc = cv2.VideoWriter_fourcc(*"avc1")
+        self.writer = cv2.VideoWriter(
+            output_path, fourcc, fps, (display_width, display_height)
+        )
+
+        # Check if writer opened successfully, fallback to mp4v if not
+        if not self.writer.isOpened():
+            fourcc = cv2.VideoWriter_fourcc(*"mp4v")
+            self.writer = cv2.VideoWriter(
+                output_path, fourcc, fps, (display_width, display_height)
+            )
+
+        if not self.writer.isOpened():
+            raise ValueError(
+                f"Failed to create video writer for {output_path} with dimensions "
+                f"{display_width}x{display_height}"
+            )
+
+    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 _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 = [
+                (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
+                # Additional segments for better visualization
+                (f"{side_prefix}shoulder", "nose", (150, 150, 255), 2),  # Neck
+            ]
+
+            # Draw ALL visible segments (not just one side)
+            for start_key, end_key, color, thickness in segments:
+                if start_key in landmarks and end_key in landmarks:
+                    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
+                        )
+
+            # Draw joints as circles for this side
+            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)
+
+        # 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)
+        if angles.get("ankle_angle") is not None:
+            self._draw_angle_arc(
+                frame, landmarks, f"{side_used}_ankle", angles["ankle_angle"]
+            )
+        if angles.get("knee_angle") is not None:
+            self._draw_angle_arc(
+                frame, landmarks, f"{side_used}_knee", angles["knee_angle"]
+            )
+        if angles.get("hip_angle") is not None:
+            self._draw_angle_arc(
+                frame, landmarks, f"{side_used}_hip", angles["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 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:
+            # Draw skeleton segments for triple extension
+            self._draw_skeleton(annotated, landmarks)
+
+            # Draw joint angles
+            self._draw_joint_angles(annotated, landmarks, phase_color)
+
+            # Draw foot landmarks
+            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(annotated, (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(annotated, (avg_x, avg_y), 12, phase_color, -1)
+                cv2.circle(annotated, (avg_x, avg_y), 14, (255, 255, 255), 2)
+
+        # Draw phase indicator banner
+        if phase:
+            # Phase name with background
+            phase_text = f"Phase: {phase.upper()}"
+            text_size = cv2.getTextSize(phase_text, cv2.FONT_HERSHEY_SIMPLEX, 1, 2)[0]
+            cv2.rectangle(annotated, (5, 5), (text_size[0] + 15, 45), phase_color, -1)
+            cv2.putText(
+                annotated,
+                phase_text,
+                (10, 35),
+                cv2.FONT_HERSHEY_SIMPLEX,
+                1,
+                (0, 0, 0),
+                2,
+            )
+
+        # 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
+        if metrics:
+            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(
+                    annotated,
+                    marker,
+                    (10, y_offset),
+                    cv2.FONT_HERSHEY_SIMPLEX,
+                    0.7,
+                    (255, 255, 0),
+                    2,
+                )
+                y_offset += 35
+
+            # Draw metrics summary in bottom right (last 30 frames)
+            total_frames = int(metrics.landing_frame) + 30
+            if frame_idx >= total_frames - 30:
+                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(
+                    annotated,
+                    (self.width - 320, self.height - box_height - 10),
+                    (self.width - 10, self.height - 10),
+                    (0, 0, 0),
+                    -1,
+                )
+                cv2.rectangle(
+                    annotated,
+                    (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(
+                        annotated,
+                        text,
+                        (self.width - 310, text_y),
+                        cv2.FONT_HERSHEY_SIMPLEX,
+                        0.6,
+                        (255, 255, 255),
+                        1,
+                    )
+                    text_y += 30
+
+        return annotated
+
+    def write_frame(self, frame: np.ndarray) -> None:
+        """
+        Write frame to output video.
+
+        Args:
+            frame: Video frame with shape (height, width, 3)
+
+        Raises:
+            ValueError: If frame dimensions don't match expected encoded dimensions
+        """
+        # Validate frame dimensions match expected encoded dimensions
+        frame_height, frame_width = frame.shape[:2]
+        if frame_height != self.height or frame_width != self.width:
+            raise ValueError(
+                f"Frame dimensions ({frame_width}x{frame_height}) don't match "
+                f"source dimensions ({self.width}x{self.height}). "
+                f"Aspect ratio must be preserved from source video."
+            )
+
+        # Resize to display dimensions if needed (to handle SAR)
+        if self.needs_resize:
+            frame = cv2.resize(
+                frame,
+                (self.display_width, self.display_height),
+                interpolation=cv2.INTER_LANCZOS4,
+            )
+
+        self.writer.write(frame)
+
+    def close(self) -> None:
+        """Release video writer."""
+        self.writer.release()
+
+    def __enter__(self) -> "CMJDebugOverlayRenderer":
+        return self
+
+    def __exit__(self, exc_type: type, exc_val: Exception, exc_tb: object) -> None:
+        self.close()