84 lines
2.5 KiB
Python
84 lines
2.5 KiB
Python
import os
|
|
import subprocess
|
|
import cv2
|
|
import mediapipe as mp
|
|
|
|
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 beep():
|
|
subprocess.Popen(['afplay', BEEP_FILE],
|
|
stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
|
|
|
|
|
|
cap = cv2.VideoCapture(0)
|
|
count = 0
|
|
stage = None
|
|
mode = 'p' # default: push-ups
|
|
|
|
HINT = '[P]ush [U]pull [B]ench [C]url [S]itup [L]plank [R]eset [Q]uit'
|
|
|
|
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)
|
|
|
|
try:
|
|
lms = results.pose_landmarks.landmark
|
|
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
|
|
|
|
# --- UI ---
|
|
h, w = image.shape[:2]
|
|
ex = EXERCISES[mode]
|
|
color = ex['color']
|
|
|
|
cv2.rectangle(image, (0, 0), (w, 80), (15, 15, 15), -1)
|
|
cv2.putText(image, ex['name'], (12, 30),
|
|
cv2.FONT_HERSHEY_SIMPLEX, 0.85, color, 2)
|
|
cv2.putText(image, f'{ex["unit"]}: {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)
|
|
|
|
cv2.rectangle(image, (0, h - 30), (w, h), (15, 15, 15), -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)
|
|
|
|
cv2.imshow('Exercise Counter', image)
|
|
|
|
key = cv2.waitKey(10) & 0xFF
|
|
if key == ord('q'):
|
|
break
|
|
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()
|