File size: 8,038 Bytes
9cbbabb
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
import mediapipe as mp
import time
from pose_estimation.angle_calculation import calculate_angle

class PushUp:
    def __init__(self):
        self.counter = 0
        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_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)

        # Stage and Counter Logic
        current_angle_for_logic = angle_left
        current_time = time.time()

        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
            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
    
    # Inside the PushUp class in push_up.py
    def reset_reps(self):
        self.counter = 0
        self.stage = "up"  # Reset to the initial stage
        self.last_counter_update = time.time() # Important to reset debounce timer
        print("PushUp reps and stage reset for new set.")