jon/sports-webcam
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Ported to ultralytics for better POSE estimation

Browse Source
This commit is contained in:
Jon
2026-05-13 17:07:43 +01:00
parent 3a57de19b2
commit 2a23eefaa3
1 changed files with 53 additions and 48 deletions
Download Patch File Download Diff File Expand all files Collapse all files

101
main.py
View File

@@ -1,43 +1,51 @@
import os import os
import subprocess import subprocess
import cv2 import cv2
import mediapipe as mp from ultralytics import YOLO
from exercises import EXERCISES from exercises import EXERCISES
mp_pose = mp.solutions.pose BEEP_FILE = os.path.join(os.path.dirname(__file__), 'beep.mp3')
mp_draw = mp.solutions.drawing_utils 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'
BEEP_FILE = os.path.join(os.path.dirname(__file__), 'beep.mp3')
def beep(): def beep():
subprocess.Popen(['afplay', BEEP_FILE], subprocess.Popen(['afplay', BEEP_FILE],
stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL) stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
cap = cv2.VideoCapture(0) model = YOLO(MODEL_PATH)
cap = cv2.VideoCapture(0)
count = 0 count = 0
stage = None stage = None
mode = 'p' # default: push-ups mode = 'p'
HINT = '[P]ush [U]pull [B]ench [C]url [S]itup [L]plank [R]eset [Q]uit' while cap.isOpened():
ret, frame = cap.read()
if not ret:
break
with mp_pose.Pose(min_detection_confidence=0.6, min_tracking_confidence=0.6) as pose: results = model(frame, verbose=False)
while cap.isOpened(): result = results[0]
ret, frame = cap.read()
if not ret:
break
image = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) # Draw YOLO skeleton on a copy of the frame (no bounding boxes or conf labels)
image.flags.writeable = False image = result.plot(boxes=False, conf=False, labels=False)
results = pose.process(image)
image.flags.writeable = True
image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
# 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: try:
lms = results.pose_landmarks.landmark ex = EXERCISES[mode]
ex = EXERCISES[mode] new_stage, counted = ex['fn'](kps, conf, stage)
new_stage, counted = ex['fn'](lms, stage)
if new_stage != stage: if new_stage != stage:
stage = new_stage stage = new_stage
beep() beep()
@@ -46,38 +54,35 @@ with mp_pose.Pose(min_detection_confidence=0.6, min_tracking_confidence=0.6) as
except Exception: except Exception:
pass pass
# --- UI --- # ── UI overlay ───────────────────────────────────────────────────────────
h, w = image.shape[:2] h, w = image.shape[:2]
ex = EXERCISES[mode] ex = EXERCISES[mode]
color = ex['color'] color = ex['color']
cv2.rectangle(image, (0, 0), (w, 80), (15, 15, 15), -1) cv2.rectangle(image, (0, 0), (w, 80), (15, 15, 15), -1)
cv2.putText(image, ex['name'], (12, 30), cv2.putText(image, ex['name'], (12, 30),
cv2.FONT_HERSHEY_SIMPLEX, 0.85, color, 2) cv2.FONT_HERSHEY_SIMPLEX, 0.85, color, 2)
cv2.putText(image, f'{ex["unit"]}: {count}', (12, 68), cv2.putText(image, f'{ex["unit"]}: {count}', (12, 68),
cv2.FONT_HERSHEY_SIMPLEX, 1.3, (0, 255, 0), 3) cv2.FONT_HERSHEY_SIMPLEX, 1.3, (0, 255, 0), 3)
cv2.putText(image, (stage or '---').upper(), (w - 108, 50), cv2.putText(image, (stage or '---').upper(), (w - 108, 50),
cv2.FONT_HERSHEY_SIMPLEX, 0.9, (180, 180, 0), 2) cv2.FONT_HERSHEY_SIMPLEX, 0.9, (180, 180, 0), 2)
cv2.rectangle(image, (0, h - 30), (w, h), (15, 15, 15), -1) cv2.rectangle(image, (0, h - 30), (w, h), (15, 15, 15), -1)
cv2.putText(image, HINT, (8, h - 8), cv2.putText(image, HINT, (8, h - 8),
cv2.FONT_HERSHEY_SIMPLEX, 0.48, (160, 160, 160), 1) cv2.FONT_HERSHEY_SIMPLEX, 0.48, (160, 160, 160), 1)
if results.pose_landmarks: cv2.imshow('Exercise Counter', image)
mp_draw.draw_landmarks(image, results.pose_landmarks, mp_pose.POSE_CONNECTIONS)
cv2.imshow('Exercise Counter', image) key = cv2.waitKey(1) & 0xFF
if key == ord('q'):
key = cv2.waitKey(10) & 0xFF break
if key == ord('q'): elif key == ord('r'):
break count = 0
elif key == ord('r'): stage = None
count = 0 elif key != 255 and chr(key) in EXERCISES:
stage = None mode = chr(key)
elif key != 255 and chr(key) in EXERCISES: stage = None
mode = chr(key) count = 0
stage = None
count = 0
cap.release() cap.release()
cv2.destroyAllWindows() cv2.destroyAllWindows()