Update exercises/push_up.py

exercises/push_up.py

@@ -1,77 +1,130 @@
-import cv2
+import mediapipe as mp
 import time
 from pose_estimation.angle_calculation import calculate_angle
 
 class PushUp:
     def __init__(self):
         self.counter = 0
-        self.stage = "Initial"  # 'up' or 'down'
-        self.angle_threshold_up = 150  # Upper threshold for 'up' stage
-        self.angle_threshold_down = 70  # Lower threshold for 'down' stage
-        self.last_counter_update = time.time()  # Track the time of the last counter update
+        self.stage = "up"  # Changed from "Initial"
+        self.angle_threshold_up = 150
+        self.angle_threshold_down = 70
+        self.last_counter_update = time.time()
+        self.mp_pose = mp.solutions.pose # Added
 
     def calculate_shoulder_elbow_wrist_angle(self, shoulder, elbow, wrist):
         """Calculate the angle between shoulder, elbow, and wrist."""
         return calculate_angle(shoulder, elbow, wrist)
 
-    def track_push_up(self, landmarks, frame):
-        # Right side landmarks (shoulder, elbow, wrist)
-        shoulder_left = [int(landmarks[11].x * frame.shape[1]), int(landmarks[11].y * frame.shape[0])]
-        elbow_left = [int(landmarks[13].x * frame.shape[1]), int(landmarks[13].y * frame.shape[0])]
-        wrist_left = [int(landmarks[15].x * frame.shape[1]), int(landmarks[15].y * frame.shape[0])]
-
-        shoulder_right = [int(landmarks[12].x * frame.shape[1]), int(landmarks[12].y * frame.shape[0])]
-        elbow_right = [int(landmarks[14].x * frame.shape[1]), int(landmarks[14].y * frame.shape[0])]
-        wrist_right = [int(landmarks[16].x * frame.shape[1]), int(landmarks[16].y * frame.shape[0])]
-
-        # Calculate angles for push-up tracking
+    def track_push_up(self, landmarks_mp, frame_width, frame_height):
+        lm = landmarks_mp # shortcut
+
+        # Left side landmarks
+        shoulder_left = [int(lm[self.mp_pose.PoseLandmark.LEFT_SHOULDER.value].x * frame_width),
+                         int(lm[self.mp_pose.PoseLandmark.LEFT_SHOULDER.value].y * frame_height)]
+        elbow_left = [int(lm[self.mp_pose.PoseLandmark.LEFT_ELBOW.value].x * frame_width),
+                      int(lm[self.mp_pose.PoseLandmark.LEFT_ELBOW.value].y * frame_height)]
+        wrist_left = [int(lm[self.mp_pose.PoseLandmark.LEFT_WRIST.value].x * frame_width),
+                      int(lm[self.mp_pose.PoseLandmark.LEFT_WRIST.value].y * frame_height)]
+
+        # Right side landmarks
+        shoulder_right = [int(lm[self.mp_pose.PoseLandmark.RIGHT_SHOULDER.value].x * frame_width),
+                          int(lm[self.mp_pose.PoseLandmark.RIGHT_SHOULDER.value].y * frame_height)]
+        elbow_right = [int(lm[self.mp_pose.PoseLandmark.RIGHT_ELBOW.value].x * frame_width),
+                       int(lm[self.mp_pose.PoseLandmark.RIGHT_ELBOW.value].y * frame_height)]
+        wrist_right = [int(lm[self.mp_pose.PoseLandmark.RIGHT_WRIST.value].x * frame_width),
+                       int(lm[self.mp_pose.PoseLandmark.RIGHT_WRIST.value].y * frame_height)]
+
+        # Calculate angles
         angle_left = self.calculate_shoulder_elbow_wrist_angle(shoulder_left, elbow_left, wrist_left)
         angle_right = self.calculate_shoulder_elbow_wrist_angle(shoulder_right, elbow_right, wrist_right)
 
-        # Draw lines with improved style
-        self.draw_line_with_style(frame, shoulder_left, elbow_left, (0, 0, 255), 2)
-        self.draw_line_with_style(frame, elbow_left, wrist_left, (0, 0, 255), 2)
-
-        self.draw_line_with_style(frame, shoulder_right, elbow_right, (102, 0, 0), 2)
-        self.draw_line_with_style(frame, elbow_right, wrist_right, (102, 0, 0), 2)
-
-        # Draw circles to highlight key points
-        self.draw_circle(frame, shoulder_left, (0, 0, 255), 8)
-        self.draw_circle(frame, elbow_left, (0, 0, 255), 8)
-        self.draw_circle(frame, wrist_left, (0, 0, 255), 8)
-
-        self.draw_circle(frame, shoulder_right, (102, 0, 0), 8)
-        self.draw_circle(frame, elbow_right, (102, 0, 0), 8)
-        self.draw_circle(frame, wrist_right, (102, 0, 0), 8)
-
-        # Update angle text positions and display
-        angle_text_position_left = (elbow_left[0] + 10, elbow_left[1] - 10)
-        cv2.putText(frame, f'Angle: {int(angle_left)}', angle_text_position_left, cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 2)
-
-        angle_text_position_right = (elbow_right[0] + 10, elbow_right[1] - 10)
-        cv2.putText(frame, f'Angle: {int(angle_right)}', angle_text_position_right, cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 2)
-
-        # Get current time
+        # Stage and Counter Logic
+        current_angle_for_logic = angle_left 
         current_time = time.time()
 
-        # Update stage and counter
-        if angle_left > self.angle_threshold_up:
-            self.stage = "Starting position"
-        elif self.angle_threshold_down < angle_left < self.angle_threshold_up and self.stage == "Starting position":
-            self.stage = "Descent"
-        elif angle_left < self.angle_threshold_down and self.stage == "Descent":
-            self.stage = "Ascent"
-            # Increment counter only if enough time has passed since last update
-            if current_time - self.last_counter_update > 1:  # 1 second threshold
+        if current_angle_for_logic > self.angle_threshold_up:
+            self.stage = "up"
+        elif self.angle_threshold_down < current_angle_for_logic < self.angle_threshold_up and self.stage == "up":
+            self.stage = "down"
+        elif current_angle_for_logic < self.angle_threshold_down and self.stage == "down":
+            if current_time - self.last_counter_update > 1:  # 1 second debounce
                 self.counter += 1
                 self.last_counter_update = current_time
-
-        return self.counter, angle_left, self.stage
-
-    def draw_line_with_style(self, frame, start_point, end_point, color, thickness):
-        """Draw a line with specified style."""
-        cv2.line(frame, start_point, end_point, color, thickness, lineType=cv2.LINE_AA)
-
-    def draw_circle(self, frame, center, color, radius):
-        """Draw a circle with specified style."""
-        cv2.circle(frame, center, radius, color, -1)  # -1 to fill the circle
+            self.stage = "up" # Transition back to up
+            
+        feedback = self._get_push_up_feedback(angle_left, angle_right, self.stage)
+
+        return {
+            "counter": self.counter,
+            "stage": self.stage,
+            "angle_left": angle_left,
+            "angle_right": angle_right,
+            "feedback": feedback
+        }
+
+    def _get_push_up_feedback(self, angle_left, angle_right, stage):
+        feedback = "Keep going!" # Default
+        
+        if stage == "down":
+            if min(angle_left, angle_right) < self.angle_threshold_down - 5:
+                feedback = "Good depth!"
+            elif min(angle_left, angle_right) > self.angle_threshold_down + 10:
+                feedback = "Go lower."
+        elif stage == "up":
+             feedback = "Push up!" # Or "Ready"
+        
+        if abs(angle_left - angle_right) > 25:
+            feedback += " Try to keep your push-up even." if feedback != "Keep going!" else "Try to keep your push-up even."
+        
+        return feedback.strip()
+
+    def get_drawing_annotations(self, landmarks_mp, frame_width, frame_height, exercise_data_dict):
+        annotations = []
+        lm = landmarks_mp # shortcut
+
+        # Pixel coordinates
+        shoulder_left = [int(lm[self.mp_pose.PoseLandmark.LEFT_SHOULDER.value].x * frame_width),
+                         int(lm[self.mp_pose.PoseLandmark.LEFT_SHOULDER.value].y * frame_height)]
+        elbow_left = [int(lm[self.mp_pose.PoseLandmark.LEFT_ELBOW.value].x * frame_width),
+                      int(lm[self.mp_pose.PoseLandmark.LEFT_ELBOW.value].y * frame_height)]
+        wrist_left = [int(lm[self.mp_pose.PoseLandmark.LEFT_WRIST.value].x * frame_width),
+                      int(lm[self.mp_pose.PoseLandmark.LEFT_WRIST.value].y * frame_height)]
+
+        shoulder_right = [int(lm[self.mp_pose.PoseLandmark.RIGHT_SHOULDER.value].x * frame_width),
+                          int(lm[self.mp_pose.PoseLandmark.RIGHT_SHOULDER.value].y * frame_height)]
+        elbow_right = [int(lm[self.mp_pose.PoseLandmark.RIGHT_ELBOW.value].x * frame_width),
+                       int(lm[self.mp_pose.PoseLandmark.RIGHT_ELBOW.value].y * frame_height)]
+        wrist_right = [int(lm[self.mp_pose.PoseLandmark.RIGHT_WRIST.value].x * frame_width),
+                       int(lm[self.mp_pose.PoseLandmark.RIGHT_WRIST.value].y * frame_height)]
+
+        # Lines (original colors: left (0,0,255) -> BGR [255,0,0], right (102,0,0) -> BGR [0,0,102])
+        annotations.append({"type": "line", "start_point": shoulder_left, "end_point": elbow_left, "color_bgr": [255, 0, 0], "thickness": 2})
+        annotations.append({"type": "line", "start_point": elbow_left, "end_point": wrist_left, "color_bgr": [255, 0, 0], "thickness": 2})
+        annotations.append({"type": "line", "start_point": shoulder_right, "end_point": elbow_right, "color_bgr": [0, 0, 102], "thickness": 2})
+        annotations.append({"type": "line", "start_point": elbow_right, "end_point": wrist_right, "color_bgr": [0, 0, 102], "thickness": 2})
+
+        # Circles
+        annotations.append({"type": "circle", "center_point": shoulder_left, "radius": 8, "color_bgr": [255, 0, 0], "filled": True})
+        annotations.append({"type": "circle", "center_point": elbow_left, "radius": 8, "color_bgr": [255, 0, 0], "filled": True})
+        annotations.append({"type": "circle", "center_point": wrist_left, "radius": 8, "color_bgr": [255, 0, 0], "filled": True})
+        annotations.append({"type": "circle", "center_point": shoulder_right, "radius": 8, "color_bgr": [0, 0, 102], "filled": True})
+        annotations.append({"type": "circle", "center_point": elbow_right, "radius": 8, "color_bgr": [0, 0, 102], "filled": True})
+        annotations.append({"type": "circle", "center_point": wrist_right, "radius": 8, "color_bgr": [0, 0, 102], "filled": True})
+
+        # Text for angles
+        if 'angle_left' in exercise_data_dict:
+            annotations.append({"type": "text", "text_content": f"Angle L: {int(exercise_data_dict['angle_left'])}", 
+                                "position": [elbow_left[0] + 10, elbow_left[1] - 10], 
+                                "font_scale": 0.5, "color_bgr": [255, 255, 255], "thickness": 2})
+        if 'angle_right' in exercise_data_dict:
+            annotations.append({"type": "text", "text_content": f"Angle R: {int(exercise_data_dict['angle_right'])}", 
+                                "position": [elbow_right[0] + 10, elbow_right[1] - 10], 
+                                "font_scale": 0.5, "color_bgr": [255, 255, 255], "thickness": 2})
+
+        # Display main feedback from exercise_data_dict
+        if 'feedback' in exercise_data_dict:
+             annotations.append({"type": "text", "text_content": exercise_data_dict['feedback'],
+                                "position": [frame_width // 2 - 150, frame_height - 40], # Adjusted for longer text
+                                "font_scale": 0.7, "color_bgr": [0, 255, 0], "thickness": 2}) # Green for feedback
+        
+        return annotations