exercise_classifier.py

import time
from enum import Enum
from collections import deque
from typing import Tuple

from pose_estimator import LandmarkData
from angle_utils import calculate_angle, angle_to_horizontal, midpoint
import config


class Exercise(Enum):
    NONE = "none"
    SQUAT = "squat"
    PUSHUP = "pushup"
    PLANK = "plank"


class ExerciseState(Enum):
    IDLE = "idle"
    REP_DOWN = "rep_down"
    REP_UP = "rep_up"
    HOLD = "hold"


class ExerciseClassifier:
    def __init__(self, history_size: int = 30):
        self.history_size = history_size
        self.knee_angle_history: deque = deque(maxlen=history_size)
        self.elbow_angle_history: deque = deque(maxlen=history_size)
        self.landmark_history: deque = deque(maxlen=history_size)

        self.current_exercise = Exercise.NONE
        self.current_state = ExerciseState.IDLE
        self.rep_count = 0
        self.hold_start_time: float = 0.0
        self.hold_time: float = 0.0

        self._static_frame_count = 0
        self._lock_frames = 0

    def classify(
        self, landmark_data: LandmarkData
    ) -> Tuple[Exercise, ExerciseState, int, float]:
        """Classify exercise and track reps.
        Returns (exercise, state, rep_count, hold_time)."""
        coords = landmark_data.normalized_coords
        vis = landmark_data.visibility

        # Need key landmarks to be visible
        required = [11, 12, 23, 24]
        if not all(idx in coords and vis.get(idx, 0) > 0.5 for idx in required):
            return self.current_exercise, self.current_state, self.rep_count, self.hold_time

        # Body orientation
        mid_shoulder = midpoint(coords[11], coords[12])
        mid_hip = midpoint(coords[23], coords[24])
        orientation_angle = angle_to_horizontal(mid_shoulder, mid_hip)

        # Compute key angles
        knee_angle = self._avg_angle(coords, vis, "knee")
        elbow_angle = self._avg_angle(coords, vis, "elbow")

        self.knee_angle_history.append(knee_angle)
        self.elbow_angle_history.append(elbow_angle)
        self.landmark_history.append(coords)

        # Decrement exercise lock
        if self._lock_frames > 0:
            self._lock_frames -= 1

        # Classification
        if self._lock_frames <= 0:
            detected = self._detect_exercise(orientation_angle, knee_angle, elbow_angle)
            if detected != self.current_exercise:
                self.current_exercise = detected
                self.current_state = ExerciseState.IDLE
                self.rep_count = 0
                self.hold_start_time = 0.0
                self.hold_time = 0.0
                self._static_frame_count = 0
                self._lock_frames = config.EXERCISE_LOCK_FRAMES

        # Update state machine
        self._update_state(knee_angle, elbow_angle)

        return self.current_exercise, self.current_state, self.rep_count, self.hold_time

    def _avg_angle(
        self, coords: dict, vis: dict, joint: str
    ) -> float:
        """Compute average angle from both sides for the given joint."""
        if joint == "knee":
            # hip-knee-ankle
            left_ok = all(vis.get(i, 0) > 0.5 for i in [23, 25, 27])
            right_ok = all(vis.get(i, 0) > 0.5 for i in [24, 26, 28])
            angles = []
            if left_ok:
                angles.append(calculate_angle(coords[23], coords[25], coords[27]))
            if right_ok:
                angles.append(calculate_angle(coords[24], coords[26], coords[28]))
            return sum(angles) / len(angles) if angles else 180.0
        elif joint == "elbow":
            # shoulder-elbow-wrist
            left_ok = all(vis.get(i, 0) > 0.5 for i in [11, 13, 15])
            right_ok = all(vis.get(i, 0) > 0.5 for i in [12, 14, 16])
            angles = []
            if left_ok:
                angles.append(calculate_angle(coords[11], coords[13], coords[15]))
            if right_ok:
                angles.append(calculate_angle(coords[12], coords[14], coords[16]))
            return sum(angles) / len(angles) if angles else 180.0
        return 180.0

    def _detect_exercise(
        self, orientation_angle: float, knee_angle: float, elbow_angle: float
    ) -> Exercise:
        """Determine which exercise based on orientation and angles."""
        if orientation_angle >= config.UPRIGHT_ANGLE_MIN:
            # Upright -> check for squat
            if knee_angle < config.SQUAT_KNEE_FLEXED_MAX + 20:
                return Exercise.SQUAT
            if len(self.knee_angle_history) >= 10:
                min_k = min(list(self.knee_angle_history)[-10:])
                if min_k < config.SQUAT_KNEE_FLEXED_MAX + 10:
                    return Exercise.SQUAT
            return Exercise.NONE

        elif orientation_angle <= config.HORIZONTAL_ANGLE_MAX:
            # Horizontal -> push-up or plank
            variance = self._landmark_variance()

            if elbow_angle < config.PUSHUP_ELBOW_FLEXED_MAX + 20:
                self._static_frame_count = 0
                return Exercise.PUSHUP

            if len(self.elbow_angle_history) >= 10:
                min_e = min(list(self.elbow_angle_history)[-10:])
                max_e = max(list(self.elbow_angle_history)[-10:])
                if max_e - min_e > 30:
                    self._static_frame_count = 0
                    return Exercise.PUSHUP

            if variance < config.PLANK_VARIANCE_THRESHOLD:
                self._static_frame_count += 1
            else:
                self._static_frame_count = 0

            if self._static_frame_count >= config.PLANK_MIN_HOLD_FRAMES:
                return Exercise.PLANK

            # Default horizontal but no clear pattern yet
            if self.current_exercise in (Exercise.PUSHUP, Exercise.PLANK):
                return self.current_exercise
            return Exercise.NONE

        return self.current_exercise if self.current_exercise != Exercise.NONE else Exercise.NONE

    def _landmark_variance(self) -> float:
        """Average movement of key landmarks between last two frames."""
        if len(self.landmark_history) < 2:
            return 1.0
        prev = self.landmark_history[-2]
        curr = self.landmark_history[-1]
        key_indices = [11, 12, 23, 24, 25, 26, 27, 28]
        total = 0.0
        count = 0
        for idx in key_indices:
            if idx in prev and idx in curr:
                dx = curr[idx][0] - prev[idx][0]
                dy = curr[idx][1] - prev[idx][1]
                total += (dx ** 2 + dy ** 2) ** 0.5
                count += 1
        return total / count if count > 0 else 1.0

    def _update_state(self, knee_angle: float, elbow_angle: float):
        """Update rep counting state machine."""
        if self.current_exercise == Exercise.SQUAT:
            if self.current_state == ExerciseState.IDLE:
                if knee_angle < config.SQUAT_KNEE_FLEXED_MAX:
                    self.current_state = ExerciseState.REP_DOWN
            elif self.current_state == ExerciseState.REP_DOWN:
                if knee_angle > config.SQUAT_KNEE_EXTENDED_MIN:
                    self.current_state = ExerciseState.REP_UP
                    self.rep_count += 1
            elif self.current_state == ExerciseState.REP_UP:
                if knee_angle < config.SQUAT_KNEE_FLEXED_MAX:
                    self.current_state = ExerciseState.REP_DOWN

        elif self.current_exercise == Exercise.PUSHUP:
            if self.current_state == ExerciseState.IDLE:
                if elbow_angle < config.PUSHUP_ELBOW_FLEXED_MAX:
                    self.current_state = ExerciseState.REP_DOWN
            elif self.current_state == ExerciseState.REP_DOWN:
                if elbow_angle > config.PUSHUP_ELBOW_EXTENDED_MIN:
                    self.current_state = ExerciseState.REP_UP
                    self.rep_count += 1
            elif self.current_state == ExerciseState.REP_UP:
                if elbow_angle < config.PUSHUP_ELBOW_FLEXED_MAX:
                    self.current_state = ExerciseState.REP_DOWN

        elif self.current_exercise == Exercise.PLANK:
            if self.current_state != ExerciseState.HOLD:
                self.current_state = ExerciseState.HOLD
                self.hold_start_time = time.time()
            self.hold_time = time.time() - self.hold_start_time

    def reset(self):
        """Reset all state."""
        self.knee_angle_history.clear()
        self.elbow_angle_history.clear()
        self.landmark_history.clear()
        self.current_exercise = Exercise.NONE
        self.current_state = ExerciseState.IDLE
        self.rep_count = 0
        self.hold_start_time = 0.0
        self.hold_time = 0.0
        self._static_frame_count = 0
        self._lock_frames = 0