exercises/hammer_curl.py

import mediapipe as mp
from pose_estimation.angle_calculation import calculate_angle

class HammerCurl:
    def __init__(self):
        self.counter_right = 0
        self.counter_left = 0
        self.stage_right = "down"
        self.stage_left = "down"
        self.mp_pose = mp.solutions.pose

        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 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 (for curl).
        # elbow is the vertex.
        return calculate_angle(shoulder, elbow, wrist)

    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