Update exercises/hammer_curl.py

exercises/hammer_curl.py

@@ -1,114 +1,169 @@
-import cv2
-import numpy as np
+import mediapipe as mp
 from pose_estimation.angle_calculation import calculate_angle
-from voice_feedback.feedback import  provide_hammer_curl_feedback , speak
-
 
 class HammerCurl:
     def __init__(self):
         self.counter_right = 0
         self.counter_left = 0
-        self.stage_right = None  # 'up' or 'down' for right arm
-        self.stage_left = None  # 'up' or 'down' for left arm
-
-        self.angle_threshold = 40  # Angle threshold for misalignment
-        self.flexion_angle_up = 155  # Flexion angle for 'up' stage
-        self.flexion_angle_down = 35  # Flexion angle for 'down' stage
+        self.stage_right = "down"
+        self.stage_left = "down"
+        self.mp_pose = mp.solutions.pose
 
-        self.angle_threshold_up = 155  # Upper threshold for 'up' stage
-        self.angle_threshold_down = 47  # Lower threshold for 'down' stage
+        self.angle_alignment_threshold = 40  # Threshold for shoulder-elbow-hip alignment
+        self.angle_curl_extended = 155     # Angle when arm is considered extended
+        self.angle_curl_flexed = 47        # Angle when arm is considered flexed
 
     def calculate_shoulder_elbow_hip_angle(self, shoulder, elbow, hip):
-        """Calculate the angle between shoulder, elbow, and hip."""
+        # Calculate the angle between elbow, shoulder, and hip (for alignment).
+        # elbow is the vertex.
         return calculate_angle(elbow, shoulder, hip)
 
     def calculate_shoulder_elbow_wrist(self, shoulder, elbow, wrist):
-        """Calculate the angle between shoulder, elbow, and wrist."""
+        # Calculate the angle between shoulder, elbow, and wrist (for curl).
+        # elbow is the vertex.
         return calculate_angle(shoulder, elbow, wrist)
 
-    def track_hammer_curl(self, landmarks, frame):
-        # Right arm landmarks (shoulder, elbow, hip, wrist)
-        shoulder_right = [int(landmarks[11].x * frame.shape[1]), int(landmarks[11].y * frame.shape[0])]
-        elbow_right = [int(landmarks[13].x * frame.shape[1]), int(landmarks[13].y * frame.shape[0])]
-        hip_right = [int(landmarks[23].x * frame.shape[1]), int(landmarks[23].y * frame.shape[0])]
-        wrist_right = [int(landmarks[15].x * frame.shape[1]), int(landmarks[15].y * frame.shape[0])]
-
-        # Left arm landmarks (shoulder, elbow, hip, wrist)
-        shoulder_left = [int(landmarks[12].x * frame.shape[1]), int(landmarks[12].y * frame.shape[0])]
-        elbow_left = [int(landmarks[14].x * frame.shape[1]), int(landmarks[14].y * frame.shape[0])]
-        hip_left = [int(landmarks[24].x * frame.shape[1]), int(landmarks[24].y * frame.shape[0])]
-        wrist_left = [int(landmarks[16].x * frame.shape[1]), int(landmarks[16].y * frame.shape[0])]
-
-        # Calculate the angle for counting (elbow flexion angle)
-        angle_right_counter = self.calculate_shoulder_elbow_wrist(shoulder_right, elbow_right, wrist_right)
-        angle_left_counter = self.calculate_shoulder_elbow_wrist(shoulder_left, elbow_left, wrist_left)
-
-        # Calculate the angle for the right arm (shoulder, elbow, hip)
-        angle_right = self.calculate_shoulder_elbow_hip_angle(shoulder_right, elbow_right, hip_right)
-
-        # Calculate the angle for the left arm (shoulder, elbow, hip)
-        angle_left = self.calculate_shoulder_elbow_hip_angle(shoulder_left, elbow_left, hip_left)
-
-        # Draw lines with improved style
-        self.draw_line_with_style(frame, shoulder_left, elbow_left, (0, 0, 255), 4)
-        self.draw_line_with_style(frame, elbow_left, wrist_left, (0, 0, 255), 4)
-
-        self.draw_line_with_style(frame, shoulder_right, elbow_right, (0, 0, 255), 4)
-        self.draw_line_with_style(frame, elbow_right, wrist_right, (0, 0, 255), 4)
-
-        # Add 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, (0, 0, 255), 8)
-        self.draw_circle(frame, elbow_right, (0, 0, 255), 8)
-        self.draw_circle(frame, wrist_right, (0, 0, 255), 8)
-
-        # Convert the angles to integers and update the text positions
-        angle_text_position_left = (elbow_left[0] + 10, elbow_left[1] - 10)
-        cv2.putText(frame, f'Angle: {int(angle_left_counter)}', 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_counter)}', angle_text_position_right, cv2.FONT_HERSHEY_SIMPLEX,
-                    0.5,
-                    (255, 255, 255), 2)
-
-        warning_message_right = None
-        warning_message_left = None
-
-        # Check for misalignment based on shoulder-elbow-hip angle
-        if abs(angle_right) > self.angle_threshold:
-            warning_message_right = f"Right Shoulder-Elbow-Hip Misalignment! Angle: {angle_right:.2f}°"
-        if abs(angle_left) > self.angle_threshold:
-            warning_message_left = f"Left Shoulder-Elbow-Hip Misalignment! Angle: {angle_left:.2f}°"
-
-        if angle_right_counter > self.angle_threshold_up:
-            self.stage_right = "Flex"
-        elif self.angle_threshold_down < angle_right_counter < self.angle_threshold_up and self.stage_right == "Flex":
-            self.stage_right = "Up"
-        elif angle_right_counter < self.angle_threshold_down and self.stage_right=="Up":
-            self.stage_right = "Down"
-            self.counter_right +=1
-
-        if angle_left_counter > self.angle_threshold_up:
-            self.stage_left = "Flex"
-        elif self.angle_threshold_down < angle_left_counter < self.angle_threshold_up and self.stage_left == "Flex":
-            self.stage_left = "Up"
-        elif angle_left_counter < self.angle_threshold_down and self.stage_left == "Up":
-            self.stage_left = "Down"
-            self.counter_left +=1
-
-        # Progress percentages: 1 for "up", 0 for "down"
-        progress_right = 1 if self.stage_right == "up" else 0
-        progress_left = 1 if self.stage_left == "up" else 0
-
-        return self.counter_right, angle_right_counter, self.counter_left, angle_left_counter, warning_message_right, warning_message_left, progress_right, progress_left, self.stage_right, self.stage_left
-
-    def draw_line_with_style(self, frame, start_point, end_point, color, thickness):
-        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
+    def track_hammer_curl(self, landmarks_mp, frame_width, frame_height):
+        lm = landmarks_mp # shortcut
+
+        # Right arm 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)]
+        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)]
+        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)]
+
+        # Left arm 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)]
+        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)]
+        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)]
+
+        # Calculate curl angles
+        angle_right_curl = self.calculate_shoulder_elbow_wrist(shoulder_right, elbow_right, wrist_right)
+        angle_left_curl = self.calculate_shoulder_elbow_wrist(shoulder_left, elbow_left, wrist_left)
+
+        # Calculate alignment angles
+        angle_right_alignment = self.calculate_shoulder_elbow_hip_angle(shoulder_right, elbow_right, hip_right)
+        angle_left_alignment = self.calculate_shoulder_elbow_hip_angle(shoulder_left, elbow_left, hip_left)
+
+        # Stage and Counter Logic - Right Arm
+        if angle_right_curl > self.angle_curl_extended:
+            self.stage_right = "down"
+        elif angle_right_curl < self.angle_curl_flexed and self.stage_right == "down":
+            self.stage_right = "up"
+            self.counter_right += 1
+        
+        # Stage and Counter Logic - Left Arm
+        if angle_left_curl > self.angle_curl_extended:
+            self.stage_left = "down"
+        elif angle_left_curl < self.angle_curl_flexed and self.stage_left == "down":
+            self.stage_left = "up"
+            self.counter_left += 1
+
+        feedback_data = self._get_hammer_curl_feedback(angle_left_curl, angle_right_curl, 
+                                                       angle_left_alignment, angle_right_alignment, 
+                                                       self.stage_left, self.stage_right)
+        
+        return {
+            "counter_left": self.counter_left,
+            "stage_left": self.stage_left,
+            "angle_left_curl": angle_left_curl,
+            "angle_left_alignment": angle_left_alignment,
+            "counter_right": self.counter_right,
+            "stage_right": self.stage_right,
+            "angle_right_curl": angle_right_curl,
+            "angle_right_alignment": angle_right_alignment,
+            "feedback_left": feedback_data["left_arm"],
+            "feedback_right": feedback_data["right_arm"]
+        }
+
+    def _get_hammer_curl_feedback(self, angle_left_curl, angle_right_curl, 
+                                 angle_left_alignment, angle_right_alignment, 
+                                 stage_left, stage_right):
+        msg_l = ""
+        msg_r = ""
+
+        if abs(angle_left_alignment) > self.angle_alignment_threshold:
+            msg_l += "Keep left elbow closer to body. "
+        if abs(angle_right_alignment) > self.angle_alignment_threshold:
+            msg_r += "Keep right elbow closer to body. "
+
+        if stage_left == "up":
+            if not msg_l: msg_l += "Left arm flexed. "
+        elif stage_left == "down":
+             if not msg_l: msg_l += "Left arm extended. "
+        
+        if stage_right == "up":
+            if not msg_r: msg_r += "Right arm flexed. "
+        elif stage_right == "down":
+            if not msg_r: msg_r += "Right arm extended. "
+
+        if not msg_l.strip(): # Use strip() to check if it's empty after potential spaces
+            msg_l = "Left: OK"
+        if not msg_r.strip():
+            msg_r = "Right: OK"
+            
+        return {"left_arm": msg_l.strip(), "right_arm": msg_r.strip()}
+
+    def get_drawing_annotations(self, landmarks_mp, frame_width, frame_height, exercise_data_dict):
+        annotations = []
+        lm = landmarks_mp 
+
+        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)]
+        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)]
+
+        color_bgr = [255, 0, 0] 
+        thickness = 4
+        annotations.append({"type": "line", "start_point": shoulder_left, "end_point": elbow_left, "color_bgr": color_bgr, "thickness": thickness})
+        annotations.append({"type": "line", "start_point": elbow_left, "end_point": wrist_left, "color_bgr": color_bgr, "thickness": thickness})
+        annotations.append({"type": "line", "start_point": shoulder_right, "end_point": elbow_right, "color_bgr": color_bgr, "thickness": thickness})
+        annotations.append({"type": "line", "start_point": elbow_right, "end_point": wrist_right, "color_bgr": color_bgr, "thickness": thickness})
+
+        radius = 8
+        annotations.append({"type": "circle", "center_point": shoulder_left, "radius": radius, "color_bgr": color_bgr, "filled": True})
+        # ... (add all other circle and text annotations as per the previous full version of get_drawing_annotations) ...
+        annotations.append({"type": "circle", "center_point": elbow_left, "radius": radius, "color_bgr": color_bgr, "filled": True})
+        annotations.append({"type": "circle", "center_point": wrist_left, "radius": radius, "color_bgr": color_bgr, "filled": True})
+        annotations.append({"type": "circle", "center_point": shoulder_right, "radius": radius, "color_bgr": color_bgr, "filled": True})
+        annotations.append({"type": "circle", "center_point": elbow_right, "radius": radius, "color_bgr": color_bgr, "filled": True})
+        annotations.append({"type": "circle", "center_point": wrist_right, "radius": radius, "color_bgr": color_bgr, "filled": True})
+
+        text_color_bgr = [255, 255, 255]
+        font_scale = 0.5
+        text_thickness = 2
+        if 'angle_left_curl' in exercise_data_dict:
+            annotations.append({"type": "text", "text_content": f"L Curl: {int(exercise_data_dict['angle_left_curl'])}", 
+                                "position": [elbow_left[0] + 10, elbow_left[1] - 20],
+                                "font_scale": font_scale, "color_bgr": text_color_bgr, "thickness": text_thickness})
+        if 'angle_right_curl' in exercise_data_dict:
+            annotations.append({"type": "text", "text_content": f"R Curl: {int(exercise_data_dict['angle_right_curl'])}", 
+                                "position": [elbow_right[0] + 10, elbow_right[1] - 20],
+                                "font_scale": font_scale, "color_bgr": text_color_bgr, "thickness": text_thickness})
+        
+        if 'feedback_left' in exercise_data_dict: # Check if key exists
+            annotations.append({"type": "text", "text_content": f"L: {exercise_data_dict['feedback_left']}",
+                                "position": [10, frame_height - 70 if frame_height > 100 else 30], 
+                                "font_scale": 0.6, "color_bgr": [0, 255, 0], "thickness": 2})
+        if 'feedback_right' in exercise_data_dict: # Check if key exists
+            annotations.append({"type": "text", "text_content": f"R: {exercise_data_dict['feedback_right']}",
+                                "position": [10, frame_height - 40 if frame_height > 70 else 60], 
+                                "font_scale": 0.6, "color_bgr": [0, 255, 0], "thickness": 2})
+        return annotations