Spaces:
Running
Running
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.") |