89 lines
2.6 KiB
Python
89 lines
2.6 KiB
Python
import os
|
|
import subprocess
|
|
import cv2
|
|
from ultralytics import YOLO
|
|
|
|
from exercises import EXERCISES
|
|
|
|
BEEP_FILE = os.path.join(os.path.dirname(__file__), 'beep.mp3')
|
|
MODEL_PATH = os.path.join(os.path.dirname(__file__), 'yolo11x-pose.pt')
|
|
|
|
HINT = '[P]ush [U]pull [B]ench [C]url [S]itup [L]plank [R]eset [Q]uit'
|
|
|
|
|
|
def beep():
|
|
subprocess.Popen(['afplay', BEEP_FILE],
|
|
stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
|
|
|
|
|
|
model = YOLO(MODEL_PATH)
|
|
|
|
cap = cv2.VideoCapture(0)
|
|
count = 0
|
|
stage = None
|
|
mode = 'p'
|
|
|
|
while cap.isOpened():
|
|
ret, frame = cap.read()
|
|
if not ret:
|
|
break
|
|
|
|
results = model(frame, verbose=False)
|
|
result = results[0]
|
|
|
|
# Draw YOLO skeleton on a copy of the frame (no bounding boxes or conf labels)
|
|
image = result.plot(boxes=False, conf=False, labels=False)
|
|
|
|
# Extract the first detected person's keypoints
|
|
kps = conf = None
|
|
if (result.keypoints is not None
|
|
and result.keypoints.xyn is not None
|
|
and len(result.keypoints.xyn) > 0):
|
|
kps = result.keypoints.xyn[0].cpu().numpy() # (17, 2) normalised
|
|
conf = result.keypoints.conf[0].cpu().numpy() # (17,)
|
|
|
|
if kps is not None:
|
|
try:
|
|
ex = EXERCISES[mode]
|
|
new_stage, counted = ex['fn'](kps, conf, stage)
|
|
if new_stage != stage:
|
|
stage = new_stage
|
|
beep()
|
|
if counted:
|
|
count += 1
|
|
except Exception:
|
|
pass
|
|
|
|
# ── UI overlay ───────────────────────────────────────────────────────────
|
|
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)
|
|
|
|
cv2.imshow('Exercise Counter', image)
|
|
|
|
key = cv2.waitKey(1) & 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()
|