142 lines
5.0 KiB
Python
142 lines
5.0 KiB
Python
import cv2
|
|
import mediapipe as mp
|
|
import numpy as np
|
|
|
|
mp_pose = mp.solutions.pose
|
|
mp_draw = mp.solutions.drawing_utils
|
|
|
|
|
|
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
|
|
|
|
|
|
cap = cv2.VideoCapture(0)
|
|
count = 0
|
|
stage = None # current rep phase
|
|
mode = 'pushup' # 'pushup' | 'pullup'
|
|
|
|
with mp_pose.Pose(min_detection_confidence=0.6, min_tracking_confidence=0.6) as pose:
|
|
while cap.isOpened():
|
|
ret, frame = cap.read()
|
|
if not ret:
|
|
break
|
|
|
|
image = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
|
|
image.flags.writeable = False
|
|
results = pose.process(image)
|
|
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
|
|
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)
|
|
|
|
# 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, f'REPS: {count}', (12, 68),
|
|
cv2.FONT_HERSHEY_SIMPLEX, 1.3, (0, 255, 0), 3)
|
|
|
|
# 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)
|
|
|
|
if results.pose_landmarks:
|
|
mp_draw.draw_landmarks(image, results.pose_landmarks, mp_pose.POSE_CONNECTIONS)
|
|
|
|
cv2.imshow('Exercise Counter', image)
|
|
|
|
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
|
|
|
|
cap.release()
|
|
cv2.destroyAllWindows()
|