posefit-server/app/exercises/dead_bug/detector.py

from __future__ import annotations

import threading
import time

import cv2
import mediapipe as mp
import numpy as np
from loguru import logger

from app.exercises.dead_bug.metrics import calculate_metrics
from app.exercises.dead_bug.rules import has_required_visibility
from app.exercises.dead_bug.state_machine import DeadBugStateMachine
from app.exercises.dead_bug.types import DeadBugMetrics, DeadBugPhase, DeadBugResult, Point
from app.rendering.overlay_renderer import draw_status_overlay
from app.rendering.skeleton_renderer import draw_landmarks
from app.vision.frame_utils import bgr_to_rgba
from app.vision.pose_landmarker import PoseLandmarkerWrapper
from app.vision.pose_types import (
    LEFT_ANKLE,
    LEFT_ELBOW,
    LEFT_HIP,
    LEFT_KNEE,
    LEFT_SHOULDER,
    LEFT_WRIST,
    REQUIRED_LANDMARKS,
    RIGHT_ANKLE,
    RIGHT_ELBOW,
    RIGHT_HIP,
    RIGHT_KNEE,
    RIGHT_SHOULDER,
    RIGHT_WRIST,
)


class DeadBugDetector:
    def __init__(
        self,
        *,
        model_path: str | None = None,
        visibility_threshold: float = 0.45,
        extension_confirm_frames: int = 4,
        reset_confirm_frames: int = 3,
        prefer_gpu: bool = True,
    ) -> None:
        self.visibility_threshold = visibility_threshold

        self._latest_result = None
        self._result_lock = threading.Lock()
        self._result_event = threading.Event()
        self._inflight = False
        self._inflight_started_at = 0.0

        def on_result(pose_result, _image, _timestamp_ms):
            with self._result_lock:
                self._latest_result = pose_result
                self._inflight = False
                self._inflight_started_at = 0.0
            self._result_event.set()

        self._landmarker = PoseLandmarkerWrapper(
            model_path=model_path,
            prefer_gpu=prefer_gpu,
            result_callback=on_result,
        )

        self._state = DeadBugStateMachine(
            extension_confirm_frames=extension_confirm_frames,
            reset_confirm_frames=reset_confirm_frames,
        )

        self._last_timestamp_ms = -1

    def close(self) -> None:
        self._landmarker.close()

    def process_frame(self, bgr_frame: np.ndarray, timestamp_ms: int) -> tuple[np.ndarray, DeadBugResult]:
        timestamp_ms = self._normalize_timestamp(timestamp_ms)

        with self._result_lock:
            if self._inflight and time.monotonic() - self._inflight_started_at > 0.5:
                logger.warning("MediaPipe detect_async timed out; allowing next frame submission")
                self._inflight = False
                self._inflight_started_at = 0.0
            should_submit = not self._inflight
            if should_submit:
                self._inflight = True
                self._inflight_started_at = time.monotonic()

        if should_submit:
            rgba_frame = bgr_to_rgba(bgr_frame)
            mp_image = mp.Image(image_format=mp.ImageFormat.SRGBA, data=rgba_frame)
            self._result_event.clear()
            try:
                self._landmarker.detect_async(mp_image, timestamp_ms)
            except Exception:
                with self._result_lock:
                    self._inflight = False
                    self._inflight_started_at = 0.0
                raise
            self._result_event.wait(timeout=0.08)

        with self._result_lock:
            pose_result = self._latest_result

        annotated = bgr_frame.copy()

        if pose_result is None or not pose_result.pose_landmarks:
            result = DeadBugResult(
                rep_count=self._state.rep_count,
                phase=DeadBugPhase.NO_POSE,
                side=self._state.active_side,
                is_standard=False,
                feedback=["No full body detected"],
                metrics=None,
            )
            draw_status_overlay(annotated, result)
            return annotated, result

        landmarks = [Point(lm.x, lm.y, lm.z, getattr(lm, "visibility", 1.0)) for lm in pose_result.pose_landmarks[0]]
        draw_landmarks(annotated, landmarks, REQUIRED_LANDMARKS, visibility_threshold=self.visibility_threshold)

        if not has_required_visibility(landmarks, REQUIRED_LANDMARKS, self.visibility_threshold):
            result = DeadBugResult(
                rep_count=self._state.rep_count,
                phase=DeadBugPhase.NO_POSE,
                side=self._state.active_side,
                is_standard=False,
                feedback=["Keep shoulders, elbows, wrists, hips, knees, ankles visible"],
                metrics=None,
            )
            draw_status_overlay(annotated, result)
            return annotated, result

        raw = calculate_metrics(
            landmarks,
            left_shoulder=LEFT_SHOULDER,
            right_shoulder=RIGHT_SHOULDER,
            left_elbow=LEFT_ELBOW,
            right_elbow=RIGHT_ELBOW,
            left_wrist=LEFT_WRIST,
            right_wrist=RIGHT_WRIST,
            left_hip=LEFT_HIP,
            right_hip=RIGHT_HIP,
            left_knee=LEFT_KNEE,
            right_knee=RIGHT_KNEE,
            left_ankle=LEFT_ANKLE,
            right_ankle=RIGHT_ANKLE,
            visibility_threshold=self.visibility_threshold,
        )

        metrics = DeadBugMetrics(
            left_arm_extended=raw["left_arm_extended"],
            right_arm_extended=raw["right_arm_extended"],
            left_leg_extended=raw["left_leg_extended"],
            right_leg_extended=raw["right_leg_extended"],
            left_elbow_angle=raw["left_elbow_angle"],
            right_elbow_angle=raw["right_elbow_angle"],
            left_knee_angle=raw["left_knee_angle"],
            right_knee_angle=raw["right_knee_angle"],
            feedback=raw["feedback"],
        )

        result = self._state.update(metrics)
        draw_status_overlay(annotated, result)
        return annotated, result

    def _normalize_timestamp(self, timestamp_ms: int) -> int:
        if timestamp_ms <= self._last_timestamp_ms:
            timestamp_ms = self._last_timestamp_ms + 1
        self._last_timestamp_ms = timestamp_ms
        return timestamp_ms