diff --git a/main.py b/main.py index 7d1283c..33038fa 100644 --- a/main.py +++ b/main.py @@ -1,64 +1,26 @@ +import os +import subprocess import cv2 import mediapipe as mp -import numpy as np + +from exercises import EXERCISES mp_pose = mp.solutions.pose mp_draw = mp.solutions.drawing_utils +BEEP_FILE = os.path.join(os.path.dirname(__file__), 'beep.mp3') -def calculate_angle(a, b, c): - a, b, c = np.array(a), np.array(b), np.array(c) - radians = np.arctan2(c[1] - b[1], c[0] - b[0]) - np.arctan2(a[1] - b[1], a[0] - b[0]) - angle = np.abs(radians * 180.0 / np.pi) - if angle > 180: - angle = 360 - angle - return angle - - -def elbow_angle_and_vis(landmarks, side): - P = mp_pose.PoseLandmark - if side == 'left': - sh, el, wr = P.LEFT_SHOULDER, P.LEFT_ELBOW, P.LEFT_WRIST - else: - sh, el, wr = P.RIGHT_SHOULDER, P.RIGHT_ELBOW, P.RIGHT_WRIST - sh_lm, el_lm, wr_lm = landmarks[sh.value], landmarks[el.value], landmarks[wr.value] - pts = [[sh_lm.x, sh_lm.y], [el_lm.x, el_lm.y], [wr_lm.x, wr_lm.y]] - vis = min(sh_lm.visibility, el_lm.visibility, wr_lm.visibility) - return calculate_angle(*pts), vis - - -def best_elbow_angle(landmarks): - l_ang, l_vis = elbow_angle_and_vis(landmarks, 'left') - r_ang, r_vis = elbow_angle_and_vis(landmarks, 'right') - if l_vis > 0.5 and r_vis > 0.5: - return (l_ang + r_ang) / 2, max(l_vis, r_vis) - return (l_ang, l_vis) if l_vis >= r_vis else (r_ang, r_vis) - - -def is_horizontal(landmarks): - """True when torso is roughly horizontal — filters out standing-upright false triggers.""" - P = mp_pose.PoseLandmark - nose_y = landmarks[P.NOSE.value].y - hip_y = (landmarks[P.LEFT_HIP.value].y + landmarks[P.RIGHT_HIP.value].y) / 2 - # When standing, nose is well above hips (small y vs large y in image coords). - # When horizontal (push-up), nose and hips are at similar y values. - return abs(nose_y - hip_y) < 0.38 - - -def wrists_above_shoulders(landmarks): - """True when both (or at least one) wrist is above its shoulder — pull-up hang check.""" - P = mp_pose.PoseLandmark - lms = landmarks - margin = 0.05 # must be clearly above, not just noise - l_ok = lms[P.LEFT_WRIST.value].y < lms[P.LEFT_SHOULDER.value].y - margin - r_ok = lms[P.RIGHT_WRIST.value].y < lms[P.RIGHT_SHOULDER.value].y - margin - return l_ok or r_ok +def beep(): + subprocess.Popen(['afplay', BEEP_FILE], + stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL) cap = cv2.VideoCapture(0) count = 0 -stage = None # current rep phase -mode = 'pushup' # 'pushup' | 'pullup' +stage = None +mode = 'p' # default: push-ups + +HINT = '[P]ush [U]pull [B]ench [C]url [R]eset [Q]uit' with mp_pose.Pose(min_detection_confidence=0.6, min_tracking_confidence=0.6) as pose: while cap.isOpened(): @@ -72,52 +34,34 @@ with mp_pose.Pose(min_detection_confidence=0.6, min_tracking_confidence=0.6) as image.flags.writeable = True image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR) - angle_disp = 0 - try: lms = results.pose_landmarks.landmark - angle, vis = best_elbow_angle(lms) - angle_disp = int(angle) - - if vis > 0.4: - if mode == 'pushup' and is_horizontal(lms): - if angle > 155: - stage = 'up' - if angle < 85 and stage == 'up': - stage = 'down' - count += 1 - - elif mode == 'pullup' and wrists_above_shoulders(lms): - if angle > 155: - stage = 'down' # hanging, arms straight - if angle < 90 and stage == 'down': - stage = 'up' # pulled up - count += 1 + ex = EXERCISES[mode] + new_stage, counted = ex['fn'](lms, stage) + if new_stage != stage: + stage = new_stage + beep() + if counted: + count += 1 except Exception: pass - # --- overlay UI --- - h, w = image.shape[:2] - mode_label = 'PUSH-UPS' if mode == 'pushup' else 'PULL-UPS' - mode_color = (50, 220, 100) if mode == 'pushup' else (80, 120, 255) + # --- UI --- + h, w = image.shape[:2] + ex = EXERCISES[mode] + color = ex['color'] - # top bar cv2.rectangle(image, (0, 0), (w, 80), (15, 15, 15), -1) - cv2.putText(image, mode_label, (12, 30), - cv2.FONT_HERSHEY_SIMPLEX, 0.85, mode_color, 2) + cv2.putText(image, ex['name'], (12, 30), + cv2.FONT_HERSHEY_SIMPLEX, 0.85, color, 2) cv2.putText(image, f'REPS: {count}', (12, 68), cv2.FONT_HERSHEY_SIMPLEX, 1.3, (0, 255, 0), 3) + cv2.putText(image, (stage or '---').upper(), (w - 108, 50), + cv2.FONT_HERSHEY_SIMPLEX, 0.9, (180, 180, 0), 2) - # angle + stage (top right) - cv2.putText(image, f'{angle_disp}\xb0', (w - 100, 35), - cv2.FONT_HERSHEY_SIMPLEX, 1.0, (0, 220, 220), 2) - cv2.putText(image, (stage or '---').upper(), (w - 108, 68), - cv2.FONT_HERSHEY_SIMPLEX, 0.8, (180, 180, 0), 2) - - # bottom hint bar cv2.rectangle(image, (0, h - 30), (w, h), (15, 15, 15), -1) - cv2.putText(image, '[P] Push-ups [U] Pull-ups [R] Reset [Q] Quit', - (8, h - 8), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (160, 160, 160), 1) + cv2.putText(image, HINT, (8, h - 8), + cv2.FONT_HERSHEY_SIMPLEX, 0.48, (160, 160, 160), 1) if results.pose_landmarks: mp_draw.draw_landmarks(image, results.pose_landmarks, mp_pose.POSE_CONNECTIONS) @@ -127,15 +71,13 @@ with mp_pose.Pose(min_detection_confidence=0.6, min_tracking_confidence=0.6) as key = cv2.waitKey(10) & 0xFF if key == ord('q'): break - elif key == ord('p'): - mode = 'pushup' - stage = None - elif key == ord('u'): - mode = 'pullup' - stage = None elif key == ord('r'): count = 0 stage = None + elif key != 255 and chr(key) in EXERCISES: + mode = chr(key) + stage = None + count = 0 cap.release() cv2.destroyAllWindows()