exercises/squat.py

import mediapipe as mp
from pose_estimation.angle_calculation import calculate_angle

class Squat:
    def __init__(self):
        self.counter = 0
        self.stage = "up"  # Initial stage
        self.mp_pose = mp.solutions.pose # Added for convenience

    def calculate_angle(self, point1, point2, point3): # Assuming these are pixel coordinates
        return calculate_angle(point1, point2, point3)

    def track_squat(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)]
        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_left = self.calculate_angle(shoulder_left, hip_left, knee_left)
        angle_right = self.calculate_angle(shoulder_right, hip_right, knee_right)

        # 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
        
        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})


        return annotations