Spaces:
Running
Running
File size: 8,837 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 138 139 140 141 142 143 144 145 146 147 148 |
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
# Inside the Squat class in squat.py
def reset_reps(self):
self.counter = 0
self.stage = "up" # Reset to the initial stage
print("Squat reps and stage reset for new set.") |