Update exercises/squat.py

exercises/squat.py

@@ -1,62 +1,142 @@
-import cv2
+import mediapipe as mp
 from pose_estimation.angle_calculation import calculate_angle
 
 class Squat:
     def __init__(self):
         self.counter = 0
-        self.stage = None
+        self.stage = "up"  # Initial stage
+        self.mp_pose = mp.solutions.pose # Added for convenience
 
-    def calculate_angle(self, hip, knee, ankle):
-        return calculate_angle(hip, knee, ankle)
+    def calculate_angle(self, point1, point2, point3): # Assuming these are pixel coordinates
+        return calculate_angle(point1, point2, point3)
 
-    def track_squat(self, landmarks, frame):
-        # Landmark coordinates
-        hip = [int(landmarks[23].x * frame.shape[1]), int(landmarks[23].y * frame.shape[0])]
-        knee = [int(landmarks[25].x * frame.shape[1]), int(landmarks[25].y * frame.shape[0])]
-        shoulder = [int(landmarks[11].x * frame.shape[1]), int(landmarks[11].y * frame.shape[0])]
+    def track_squat(self, landmarks_mp, frame_width, frame_height):
+        lm = landmarks_mp # shortcut
 
-        hip_right = [int(landmarks[24].x * frame.shape[1]), int(landmarks[24].y * frame.shape[0])]
-        knee_right = [int(landmarks[26].x * frame.shape[1]), int(landmarks[26].y * frame.shape[0])]
-        shoulder_right = [int(landmarks[12].x * frame.shape[1]), int(landmarks[12].y * frame.shape[0])]
+        # 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)]
+        hip_left = [int(lm[self.mp_pose.PoseLandmark.LEFT_HIP.value].x * frame_width),
+                    int(lm[self.mp_pose.PoseLandmark.LEFT_HIP.value].y * frame_height)]
+        knee_left = [int(lm[self.mp_pose.PoseLandmark.LEFT_KNEE.value].x * frame_width),
+                     int(lm[self.mp_pose.PoseLandmark.LEFT_KNEE.value].y * frame_height)]
+        # ankle_left = [int(lm[self.mp_pose.PoseLandmark.LEFT_ANKLE.value].x * frame_width),
+        #               int(lm[self.mp_pose.PoseLandmark.LEFT_ANKLE.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)]
+        hip_right = [int(lm[self.mp_pose.PoseLandmark.RIGHT_HIP.value].x * frame_width),
+                     int(lm[self.mp_pose.PoseLandmark.RIGHT_HIP.value].y * frame_height)]
+        knee_right = [int(lm[self.mp_pose.PoseLandmark.RIGHT_KNEE.value].x * frame_width),
+                      int(lm[self.mp_pose.PoseLandmark.RIGHT_KNEE.value].y * frame_height)]
+        # ankle_right = [int(lm[self.mp_pose.PoseLandmark.RIGHT_ANKLE.value].x * frame_width),
+        #                int(lm[self.mp_pose.PoseLandmark.RIGHT_ANKLE.value].y * frame_height)]
 
         # Calculate angles
-        angle = self.calculate_angle(shoulder, hip, knee)
+        angle_left = self.calculate_angle(shoulder_left, hip_left, knee_left)
         angle_right = self.calculate_angle(shoulder_right, hip_right, knee_right)
 
-        # Draw lines and circles to highlight key points
-        self.draw_line_with_style(frame, shoulder, hip, (178, 102, 255), 2)
-        self.draw_line_with_style(frame, hip, knee, (178, 102, 255), 2)
-        self.draw_line_with_style(frame, shoulder_right, hip_right, (51, 153, 255), 2)
-        self.draw_line_with_style(frame, hip_right, knee_right, (51, 153, 255), 2)
-
-        self.draw_circle(frame, shoulder, (178, 102, 255), 8)
-        self.draw_circle(frame, hip, (178, 102, 255), 8)
-        self.draw_circle(frame, knee, (178, 102, 255), 8)
-        self.draw_circle(frame, shoulder_right, (51, 153, 255), 8)
-        self.draw_circle(frame, hip_right, (51, 153, 255), 8)
-        self.draw_circle(frame, knee_right, (51, 153, 255), 8)
-
-        # Display angles on screen
-        angle_text_position = (knee[0] + 10, knee[1] - 10)
-        cv2.putText(frame, f'Angle Left: {int(angle)}', angle_text_position, cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 2)
-
-        angle_text_position_right = (knee_right[0] + 10, knee_right[1] - 10)
-        cv2.putText(frame, f'Angle Right: {int(angle_right)}', angle_text_position_right, cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 2)
-
-        # Update exercise stage and counter
-        if angle > 170:
-            self.stage = "Starting Position"
-        elif 90 < angle < 170 and self.stage == "Starting Position":
-            self.stage = "Descent"
-        elif angle < 90 and self.stage == "Descent":
-            self.stage = "Ascent"
+        # Stage and counter logic (using angle_left for primary logic)
+        current_angle_for_logic = angle_left 
+        if current_angle_for_logic > 170:
+            self.stage = "up"
+        elif 90 < current_angle_for_logic < 170 and self.stage == "up":
+            self.stage = "down"
+        elif current_angle_for_logic < 90 and self.stage == "down":
+            self.stage = "up"
             self.counter += 1
-        return self.counter, angle, self.stage
+        
+        feedback_message = self._get_squat_feedback(angle_left, angle_right, self.stage, 
+                                                    knee_left, hip_left, shoulder_left, 
+                                                    knee_right, hip_right, shoulder_right)
+
+        return {
+            "counter": self.counter,
+            "stage": self.stage,
+            "angle_left": angle_left,
+            "angle_right": angle_right,
+            "feedback": feedback_message
+        }
+
+    def _get_squat_feedback(self, angle_left, angle_right, stage, 
+                            knee_left, hip_left, shoulder_left, 
+                            knee_right, hip_right, shoulder_right): # Added points for future use
+        feedback = "Keep going." # Default feedback
+
+        if stage == "down":
+            if min(angle_left, angle_right) < 80:
+                feedback = "Good depth!"
+            elif min(angle_left, angle_right) > 100: # Knees should be more bent
+                feedback = "Go lower."
+        
+        if abs(angle_left - angle_right) > 20: # Check for uneven squat
+             # Adding a check to see if there's significant movement, e.g., not in "up" stage fully extended
+            if not (stage == "up" and min(angle_left, angle_right) > 160): # Avoid this message if standing straight
+                feedback += " Try to keep your squat even." if feedback != "Keep going." else "Try to keep your squat even."
+
+
+        # Placeholder for more advanced feedback using the passed points:
+        # E.g., Knee valgus: check if knee_left.x < hip_left.x and knee_left.x > shoulder_left.x (simplified)
+        # E.g., Back posture: calculate angle shoulder-hip-ankle (requires ankle points)
+
+        return feedback.strip()
+
+
+    def get_drawing_annotations(self, landmarks_mp, frame_width, frame_height, exercise_data_dict):
+        annotations = []
+        lm = landmarks_mp # shortcut
+
+        # Re-calculate or retrieve necessary points (pixel coordinates)
+        # For simplicity, re-calculating here. Could be optimized by passing from track_squat.
+        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)]
+        hip_left = [int(lm[self.mp_pose.PoseLandmark.LEFT_HIP.value].x * frame_width),
+                    int(lm[self.mp_pose.PoseLandmark.LEFT_HIP.value].y * frame_height)]
+        knee_left = [int(lm[self.mp_pose.PoseLandmark.LEFT_KNEE.value].x * frame_width),
+                     int(lm[self.mp_pose.PoseLandmark.LEFT_KNEE.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)]
+        hip_right = [int(lm[self.mp_pose.PoseLandmark.RIGHT_HIP.value].x * frame_width),
+                     int(lm[self.mp_pose.PoseLandmark.RIGHT_HIP.value].y * frame_height)]
+        knee_right = [int(lm[self.mp_pose.PoseLandmark.RIGHT_KNEE.value].x * frame_width),
+                      int(lm[self.mp_pose.PoseLandmark.RIGHT_KNEE.value].y * frame_height)]
+
+        # Lines for left side (original color: (178, 102, 255) -> BGR: [255, 102, 178])
+        annotations.append({"type": "line", "start_point": shoulder_left, "end_point": hip_left, "color_bgr": [255, 102, 178], "thickness": 2})
+        annotations.append({"type": "line", "start_point": hip_left, "end_point": knee_left, "color_bgr": [255, 102, 178], "thickness": 2})
+
+        # Lines for right side (original color: (51, 153, 255) -> BGR: [255, 153, 51])
+        annotations.append({"type": "line", "start_point": shoulder_right, "end_point": hip_right, "color_bgr": [255, 153, 51], "thickness": 2})
+        annotations.append({"type": "line", "start_point": hip_right, "end_point": knee_right, "color_bgr": [255, 153, 51], "thickness": 2})
+
+        # Circles for left side
+        annotations.append({"type": "circle", "center_point": shoulder_left, "radius": 8, "color_bgr": [255, 102, 178], "filled": True})
+        annotations.append({"type": "circle", "center_point": hip_left, "radius": 8, "color_bgr": [255, 102, 178], "filled": True})
+        annotations.append({"type": "circle", "center_point": knee_left, "radius": 8, "color_bgr": [255, 102, 178], "filled": True})
+
+        # Circles for right side
+        annotations.append({"type": "circle", "center_point": shoulder_right, "radius": 8, "color_bgr": [255, 153, 51], "filled": True})
+        annotations.append({"type": "circle", "center_point": hip_right, "radius": 8, "color_bgr": [255, 153, 51], "filled": True})
+        annotations.append({"type": "circle", "center_point": knee_right, "radius": 8, "color_bgr": [255, 153, 51], "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": [knee_left[0] + 10, knee_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": [knee_right[0] + 10, knee_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 - 100, frame_height - 40], # Centered at bottom
+                                "font_scale": 0.7, "color_bgr": [0, 0, 255], "thickness": 2})
 
-    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
+        return annotations