diff --git a/exercises.py b/exercises.py new file mode 100644 index 0000000..d277ce8 --- /dev/null +++ b/exercises.py @@ -0,0 +1,184 @@ +import time +import numpy as np + +# COCO keypoint indices used by YOLO pose +_NOSE = 0 +_L_SHOULDER = 5 +_R_SHOULDER = 6 +_L_ELBOW = 7 +_R_ELBOW = 8 +_L_WRIST = 9 +_R_WRIST = 10 +_L_HIP = 11 +_R_HIP = 12 +_L_ANKLE = 15 +_R_ANKLE = 16 + + +# ── helpers ────────────────────────────────────────────────────────────────── + +def _angle(a, b, c): + a, b, c = np.array(a), np.array(b), np.array(c) + rad = np.arctan2(c[1]-b[1], c[0]-b[0]) - np.arctan2(a[1]-b[1], a[0]-b[0]) + deg = np.abs(rad * 180.0 / np.pi) + return 360 - deg if deg > 180 else deg + + +def _elbow(kps, conf, side): + sh, el, wr = (_L_SHOULDER, _L_ELBOW, _L_WRIST) if side == 'left' \ + else (_R_SHOULDER, _R_ELBOW, _R_WRIST) + return _angle(kps[sh], kps[el], kps[wr]), min(conf[sh], conf[el], conf[wr]) + + +def _best_elbow(kps, conf): + la, lv = _elbow(kps, conf, 'left') + ra, rv = _elbow(kps, conf, 'right') + if lv > 0.5 and rv > 0.5: + return (la + ra) / 2, max(lv, rv) + return (la, lv) if lv >= rv else (ra, rv) + + +def _mid(kps, a, b): + return ((kps[a][0] + kps[b][0]) / 2, + (kps[a][1] + kps[b][1]) / 2) + + +def _is_horizontal(kps): + """Nose and hips at similar y → body is lying flat (side-on view).""" + nose_y = kps[_NOSE][1] + _, hip_y = _mid(kps, _L_HIP, _R_HIP) + return abs(nose_y - hip_y) < 0.38 + + +def _wrists_above_shoulders(kps, margin=0.05): + """At least one wrist clearly above its shoulder (front-on view).""" + l = kps[_L_WRIST][1] < kps[_L_SHOULDER][1] - margin + r = kps[_R_WRIST][1] < kps[_R_SHOULDER][1] - margin + return l or r + + +def _plank_deviation(kps): + """Signed hip deviation from the shoulder-to-ankle line. + 0 = straight. Positive = hips sagging. Negative = hips piking.""" + sh_x, sh_y = _mid(kps, _L_SHOULDER, _R_SHOULDER) + hi_x, hi_y = _mid(kps, _L_HIP, _R_HIP) + an_x, an_y = _mid(kps, _L_ANKLE, _R_ANKLE) + dx = an_x - sh_x + if abs(dx) < 0.01: + return 0.0 + t = (hi_x - sh_x) / dx + return hi_y - (sh_y + t * (an_y - sh_y)) + + +# ── exercise update functions ───────────────────────────────────────────────── +# All return (new_stage: str | None, rep_counted: bool) +# Side-on: push-up, curl, sit-up, plank +# Front-on: pull-up, bench + +def update_pushup(kps, conf, stage): + angle, vis = _best_elbow(kps, conf) + if vis < 0.4 or not _is_horizontal(kps): + return stage, False + if angle > 155 and stage != 'up': + return 'up', False + if angle < 85 and stage == 'up': + return 'down', True + return stage, False + + +def update_pullup(kps, conf, stage): + angle, vis = _best_elbow(kps, conf) + if vis < 0.4 or not _wrists_above_shoulders(kps): + return stage, False + if angle > 155 and stage != 'down': + return 'down', False + if angle < 90 and stage == 'down': + return 'up', True + return stage, False + + +def update_bench(kps, conf, stage): + """Horizontal body + wrists above shoulders (lying on back). Count on lockout.""" + angle, vis = _best_elbow(kps, conf) + if vis < 0.4 or not _is_horizontal(kps) or not _wrists_above_shoulders(kps): + return stage, False + if angle < 90 and stage != 'down': + return 'down', False + if angle > 155 and stage == 'down': + return 'up', True + return stage, False + + +def update_curl(kps, conf, stage): + """Standing (not horizontal). Count at full curl.""" + angle, vis = _best_elbow(kps, conf) + if vis < 0.4 or _is_horizontal(kps): + return stage, False + if angle > 150 and stage != 'down': + return 'down', False + if angle < 60 and stage == 'down': + return 'up', True + return stage, False + + +def update_situp(kps, conf, stage): + """Side-on view. Count when shoulders rise well above hips.""" + _, sh_y = _mid(kps, _L_SHOULDER, _R_SHOULDER) + _, hi_y = _mid(kps, _L_HIP, _R_HIP) + rise = hi_y - sh_y # larger = shoulders further above hips + if rise < 0.10 and stage != 'down': + return 'down', False + if rise > 0.28 and stage == 'down': + return 'up', True + return stage, False + + +class _PlankUpdater: + """Holds a per-second tick timer so the count increments in whole seconds.""" + _LIMIT = 0.08 + + def __init__(self): + self._tick = None + + def __call__(self, kps, conf, stage): + if stage is None: + self._tick = None + + dev = _plank_deviation(kps) + flat = _is_horizontal(kps) + + if not flat: + self._tick = None + return 'get flat', False + if dev > self._LIMIT: + self._tick = None + return 'sagging', False + if dev < -self._LIMIT: + self._tick = None + return 'piking', False + + now = time.time() + if self._tick is None: + self._tick = now + if now - self._tick >= 1.0: + self._tick = now + return 'holding', True + return 'holding', False + + +_plank_fn = _PlankUpdater() + +def update_plank(kps, conf, stage): + return _plank_fn(kps, conf, stage) + + +# ── registry ────────────────────────────────────────────────────────────────── + +EXERCISES = { + 'p': {'name': 'PUSH-UPS', 'fn': update_pushup, 'color': (50, 220, 100), 'unit': 'REPS'}, + 'u': {'name': 'PULL-UPS', 'fn': update_pullup, 'color': (80, 120, 255), 'unit': 'REPS'}, + 'b': {'name': 'BENCH PRESS', 'fn': update_bench, 'color': (0, 200, 255), 'unit': 'REPS'}, + 'c': {'name': 'BICEP CURLS', 'fn': update_curl, 'color': (200, 80, 255), 'unit': 'REPS'}, + 's': {'name': 'SIT-UPS', 'fn': update_situp, 'color': (255, 160, 30), 'unit': 'REPS'}, + 'l': {'name': 'PLANK', 'fn': update_plank, 'color': (30, 200, 220), 'unit': 'SECS'}, +}