scripts/app/motionbench.py

import os
import sys
import threading
import time
from pathlib import Path
from types import SimpleNamespace

import av
import cv2
import joblib
import mediapipe as mp
import numpy as np
import streamlit as st
import torch
from streamlit_webrtc import VideoProcessorBase, webrtc_streamer

PROJECT_ROOT = Path(__file__).resolve().parents[2]
if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))

from scripts.evaluate.rep_counting_methods import EXERCISE_CONFIGS, FixedThresholdFSMCounter, SmoothingBuffer, extract_primary_angle, normalize_exercise_name
from scripts.realtime_eval.evaluate_realtime_webcam import (
    MODEL_SPECS,
    build_landmark_indices,
    build_model_and_tools,
    extract_frame_features,
    get_angle_triplets,
    load_pose_module,
)


SEQUENCE_LENGTH = 30
FEATURE_COUNT = 78
DEFAULT_MODELS_ROOT = "models"
DEFAULT_PREDICTION_INTERVAL = 1.0
CAMERA_INDEX_CANDIDATES = [0, 1, 2]
RTC_CONFIGURATION = {"iceServers": [{"urls": ["stun:stun.l.google.com:19302"]}]}
BROWSER_MEDIA_CONSTRAINTS = {"video": {"width": {"ideal": 1280}, "height": {"ideal": 720}, "frameRate": {"ideal": 24, "max": 30}}, "audio": False}


def load_runtime(model_name: str, models_root: str, feature_count: int):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    args = SimpleNamespace(model_name=model_name, models_root=models_root, feature_count=feature_count)
    model, scaler, label_encoder = build_model_and_tools(args, device)
    return device, model, scaler, label_encoder


def load_similarity_asset(model_name: str, models_root: str):
    asset_path = Path(models_root) / model_name / "weights" / "similarity_centroids.pkl"
    if not asset_path.exists():
        return None
    return joblib.load(asset_path)


def cosine_similarity_percent(vector_a: np.ndarray, vector_b: np.ndarray) -> float:
    denom = float(np.linalg.norm(vector_a) * np.linalg.norm(vector_b))
    if denom <= 1e-8:
        return 0.0
    score = float(np.dot(vector_a, vector_b) / denom)
    score = max(-1.0, min(1.0, score))
    return ((score + 1.0) / 2.0) * 100.0


def read_valid_frame(capture: cv2.VideoCapture, max_reads: int = 20) -> np.ndarray | None:
    frame_bgr = None
    for _ in range(max_reads):
        ok, candidate = capture.read()
        if not ok:
            continue
        if float(np.mean(candidate)) > 5.0:
            return candidate
        frame_bgr = candidate
    if frame_bgr is not None and float(np.mean(frame_bgr)) > 5.0:
        return frame_bgr
    return None


def open_camera_with_fallback() -> cv2.VideoCapture | None:
    for camera_index in CAMERA_INDEX_CANDIDATES:
        capture = cv2.VideoCapture(camera_index)
        if not capture.isOpened():
            capture.release()
            continue
        frame_bgr = read_valid_frame(capture, max_reads=10)
        if frame_bgr is not None:
            return capture
        capture.release()
    return None


def create_runtime_state():
    return {
        "counter": None,
        "smoother": None,
        "active_exercise": None,
        "current_label": "none",
        "current_similarity": None,
        "current_reps": 0,
        "last_prediction_time": 0.0,
        "window": [],
    }


def process_single_frame(frame_bgr: np.ndarray, state: dict, model, scaler, label_encoder, device, pose_estimator, pose_module, landmark_indices, angle_triplets, similarity_asset, prediction_interval: float):
    frame_rgb = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB)
    pose = pose_estimator.process(frame_rgb)

    drawing_utils = mp.solutions.drawing_utils
    if drawing_utils is not None and pose.pose_landmarks:
        drawing_utils.draw_landmarks(frame_bgr, pose.pose_landmarks, pose_module.POSE_CONNECTIONS)

    frame_features = extract_frame_features(pose, landmark_indices, angle_triplets)
    if frame_features is not None:
        state["window"].append(frame_features)
        if len(state["window"]) > SEQUENCE_LENGTH:
            state["window"].pop(0)

    if len(state["window"]) == SEQUENCE_LENGTH and (time.time() - state["last_prediction_time"]) >= prediction_interval:
        now = time.time()
        sequence_flat = np.array(state["window"], dtype=np.float32).reshape(1, -1)
        scaled_flat = scaler.transform(sequence_flat)
        scaled = scaled_flat.reshape(1, SEQUENCE_LENGTH, FEATURE_COUNT)
        input_tensor = torch.tensor(scaled, dtype=torch.float32, device=device)
        with torch.inference_mode():
            logits = model(input_tensor)
            prediction_index = int(torch.argmax(logits, dim=1).item())
            predicted_label = label_encoder.classes_[prediction_index]
        state["current_label"] = predicted_label

        if similarity_asset is not None:
            scaled_vector = scaled_flat[0]
            centroids = similarity_asset.get("centroids", {})
            centroid_vector = centroids.get(state["current_label"])
            if centroid_vector is not None:
                state["current_similarity"] = cosine_similarity_percent(scaled_vector.astype(np.float32), np.asarray(centroid_vector, dtype=np.float32))
            else:
                state["current_similarity"] = None
        state["last_prediction_time"] = now

    normalized_label = normalize_exercise_name(state["current_label"])
    current_reps = 0
    if pose.pose_landmarks and normalized_label in EXERCISE_CONFIGS:
        if normalized_label != state["active_exercise"]:
            config = EXERCISE_CONFIGS[normalized_label]
            state["counter"] = FixedThresholdFSMCounter(config.fixed_low, config.fixed_high, config.min_state_frames)
            state["smoother"] = SmoothingBuffer(config.smoothing_window)
            state["active_exercise"] = normalized_label

        landmarks = {}
        for name, index in landmark_indices.items():
            lm = pose.pose_landmarks.landmark[index]
            landmarks[name] = np.array([lm.x, lm.y, lm.z], dtype=np.float32) if lm.visibility >= 0.5 else np.array([0.0, 0.0, 0.0], dtype=np.float32)

        config = EXERCISE_CONFIGS[normalized_label]
        raw_angle = extract_primary_angle(landmarks, config)
        smoothed_angle = state["smoother"].update(raw_angle)
        state["counter"].update(smoothed_angle)
        current_reps = state["counter"].reps
    else:
        state["active_exercise"] = None
        state["counter"] = None
        state["smoother"] = None

    state["current_reps"] = current_reps
    show_frame = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB)
    return show_frame, state


class MotionVideoProcessor(VideoProcessorBase):
    def __init__(self):
        self.lock = threading.Lock()
        self.initialized = False
        self.state = create_runtime_state()

    def configure(self, _model_name, model, scaler, label_encoder, device, pose_module, landmark_indices, angle_triplets, similarity_asset, prediction_interval):
        self.model = model
        self.scaler = scaler
        self.label_encoder = label_encoder
        self.device = device
        self.pose_module = pose_module
        self.landmark_indices = landmark_indices
        self.angle_triplets = angle_triplets
        self.similarity_asset = similarity_asset
        self.prediction_interval = prediction_interval
        self.pose_estimator = pose_module.Pose(static_image_mode=False, model_complexity=1, min_detection_confidence=0.5, min_tracking_confidence=0.5)
        self.initialized = True

    def recv(self, frame):
        frame_bgr = frame.to_ndarray(format="bgr24")
        with self.lock:
            if self.initialized:
                frame_rgb, self.state = process_single_frame(
                    frame_bgr=frame_bgr,
                    state=self.state,
                    model=self.model,
                    scaler=self.scaler,
                    label_encoder=self.label_encoder,
                    device=self.device,
                    pose_estimator=self.pose_estimator,
                    pose_module=self.pose_module,
                    landmark_indices=self.landmark_indices,
                    angle_triplets=self.angle_triplets,
                    similarity_asset=self.similarity_asset,
                    prediction_interval=self.prediction_interval,
                )
            else:
                frame_rgb = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB)
        return av.VideoFrame.from_ndarray(frame_rgb, format="rgb24")

    def __del__(self):
        if hasattr(self, "pose_estimator"):
            self.pose_estimator.close()


def render_metrics(model_name: str, state: dict, slot):
    similarity_text = f"{state['current_similarity']:0.1f}%" if state["current_similarity"] is not None else "N/A"
    slot.markdown(f"### Live Metrics\nModel: `{model_name}`\n\nExercise: `{state['current_label']}`\n\nReps: `{state['current_reps']}`\n\nSimilarity: `{similarity_text}`")


def run_local_session(model_name: str, model, scaler, label_encoder, device, pose_module, landmark_indices, angle_triplets, similarity_asset):
    capture = open_camera_with_fallback()
    if capture is None:
        st.error("No camera device found in runtime. If using Docker locally, run with --device=/dev/video0:/dev/video0. On Hugging Face Spaces, use Browser (HF/Cloud) mode.")
        st.session_state.session_active = False
        return

    left_col, right_col = st.columns([2, 1])
    with left_col:
        frame_slot = st.empty()
    with right_col:
        metrics_slot = st.empty()

    state = create_runtime_state()
    with pose_module.Pose(static_image_mode=False, model_complexity=1, min_detection_confidence=0.5, min_tracking_confidence=0.5) as pose_estimator:
        while True:
            ok, frame_bgr = capture.read()
            if not ok:
                break
            show_frame, state = process_single_frame(
                frame_bgr=frame_bgr,
                state=state,
                model=model,
                scaler=scaler,
                label_encoder=label_encoder,
                device=device,
                pose_estimator=pose_estimator,
                pose_module=pose_module,
                landmark_indices=landmark_indices,
                angle_triplets=angle_triplets,
                similarity_asset=similarity_asset,
                prediction_interval=DEFAULT_PREDICTION_INTERVAL,
            )
            frame_slot.image(show_frame, channels="RGB", width="stretch")
            render_metrics(model_name, state, metrics_slot)
            if not st.session_state.session_active:
                break

    capture.release()
    if st.session_state.session_active:
        st.session_state.session_active = False
        st.session_state.session_notice = "finished"
        st.rerun()


def run_browser_session(model_name: str, model, scaler, label_encoder, device, pose_module, landmark_indices, angle_triplets, similarity_asset):
    left_col, right_col = st.columns([2, 1])
    with left_col:
        webrtc_ctx = webrtc_streamer(
            key="motionbench-webrtc",
            video_processor_factory=MotionVideoProcessor,
            rtc_configuration=RTC_CONFIGURATION,
            media_stream_constraints=BROWSER_MEDIA_CONSTRAINTS,
            async_processing=True,
        )
    with right_col:
        metrics_slot = st.empty()

    if webrtc_ctx.video_processor:
        video_processor = webrtc_ctx.video_processor
        if not video_processor.initialized:
            video_processor.configure(model_name, model, scaler, label_encoder, device, pose_module, landmark_indices, angle_triplets, similarity_asset, DEFAULT_PREDICTION_INTERVAL)
        with video_processor.lock:
            state_copy = {
                "current_label": video_processor.state["current_label"],
                "current_reps": video_processor.state["current_reps"],
                "current_similarity": video_processor.state["current_similarity"],
            }
        render_metrics(model_name, state_copy, metrics_slot)
    else:
        metrics_slot.info("Allow camera access and click START in the video panel. If needed, choose your camera under Video Input.")


def main():
    st.set_page_config(page_title="MotionBench", layout="wide")
    st.title("MotionBench Live")

    if "session_active" not in st.session_state:
        st.session_state.session_active = False
    if "session_notice" not in st.session_state:
        st.session_state.session_notice = None

    is_hf_space = bool(os.getenv("SPACE_ID"))
    model_name = st.selectbox("Select Model", options=list(MODEL_SPECS.keys()), index=0)
    camera_options = ["Browser (HF/Cloud)", "Local OpenCV (Desktop)"]
    default_camera_index = 0 if is_hf_space else 1
    camera_source = st.selectbox("Camera Source", options=camera_options, index=default_camera_index)
    session_button_label = "Stop Session" if st.session_state.session_active else "Start Session"
    session_button_clicked = st.button(session_button_label, width="stretch")

    if session_button_clicked and st.session_state.session_active:
        st.session_state.session_active = False
        st.session_state.session_notice = "stopped"
        st.rerun()

    if session_button_clicked and not st.session_state.session_active:
        st.session_state.session_active = True
        st.session_state.session_notice = None
        st.rerun()

    if not st.session_state.session_active:
        if st.session_state.session_notice == "stopped":
            st.info("Session stopped.")
            st.session_state.session_notice = None
        elif st.session_state.session_notice == "finished":
            st.success("Session finished.")
            st.session_state.session_notice = None
        st.info("Select a model, then start session.")
        return

    device, model, scaler, label_encoder = load_runtime(model_name, DEFAULT_MODELS_ROOT, feature_count=FEATURE_COUNT)
    similarity_asset = load_similarity_asset(model_name, DEFAULT_MODELS_ROOT)
    pose_module = load_pose_module()
    landmark_indices = build_landmark_indices(pose_module)
    angle_triplets = get_angle_triplets()

    if camera_source == "Local OpenCV (Desktop)":
        run_local_session(model_name, model, scaler, label_encoder, device, pose_module, landmark_indices, angle_triplets, similarity_asset)
    else:
        run_browser_session(model_name, model, scaler, label_encoder, device, pose_module, landmark_indices, angle_triplets, similarity_asset)


if __name__ == "__main__":
    main()