u9

app.py

@@ -1,99 +1,134 @@
 import streamlit as st
-import cv2
-import mediapipe as mp
-import numpy as np
-from streamlit_webrtc import webrtc_streamer, VideoTransformerBase
+from streamlit_webrtc import webrtc_streamer, WebRtcMode, RTCConfiguration
 import av
+import numpy as np
+import mediapipe as mp
+from typing import List, Tuple, Dict
+import logging
+from PIL import Image
+import queue
 import threading
-from dataclasses import dataclass
-from typing import List
+import logging
+import sys
 
-# Mediapipe setup
-mp_drawing = mp.solutions.drawing_utils
+try:
+    from typing import Literal
+except ImportError:
+    from typing_extensions import Literal
+
+# Configure logging
+logging.basicConfig(level=logging.DEBUG)
+logger = logging.getLogger(__name__)
+
+# Initialize MediaPipe Pose
 mp_pose = mp.solutions.pose
+mp_drawing = mp.solutions.drawing_utils
 
-# Custom CSS
+# Page config
+st.set_page_config(page_title="AI Workout Trainer", page_icon="💪")
+
+# Constants
+RTC_CONFIGURATION = RTCConfiguration(
+    {"iceServers": [{"urls": ["stun:stun.l.google.com:19302"]}]}
+)
+
+# CSS
 st.markdown("""
 <style>
     .main {
-        background: linear-gradient(135deg, #001f3f 0%, #00b4d8 100%);
+        background: linear-gradient(135deg, #1e3c72 0%, #2a5298 100%);
     }
     .stButton > button {
-        background-color: #00b4d8;
+        width: 100%;
+        background-color: #2a5298;
         color: white;
-        border: none;
-        padding: 0.5rem 2rem;
         border-radius: 5px;
-        margin: 0.5rem;
-        transition: all 0.3s;
-    }
-    .stButton > button:hover {
-        background-color: #0077b6;
-    }
-    h1, h2, h3 {
-        color: #001f3f;
+        padding: 10px;
+        margin: 5px 0;
     }
     .workout-container {
-        background: rgba(0, 180, 216, 0.1);
-        padding: 2rem;
+        background-color: rgba(255, 255, 255, 0.1);
+        padding: 20px;
         border-radius: 10px;
-        margin: 1rem 0;
+        margin: 10px 0;
     }
-    .feedback-text {
-        background: rgba(0, 31, 63, 0.1);
-        padding: 1rem;
-        border-radius: 5px;
-        margin: 1rem 0;
+    h1, h2, h3 {
+        color: white;
     }
 </style>
 """, unsafe_allow_html=True)
 
-@dataclass
-class ExerciseState:
-    counter: int = 0
-    stage: str = None
-    feedback: str = ""
-
-# Global state
-state = ExerciseState()
-lock = threading.Lock()
-
-def calculate_angle(a, b, c):
-    """Calculate angle between three points."""
-    a = np.array(a)
-    b = np.array(b)
-    c = np.array(c)
-    
-    radians = np.arctan2(c[1]-b[1], c[0]-b[0]) - np.arctan2(a[1]-b[1], a[0]-b[0])
-    angle = np.abs(np.degrees(radians))
-    
-    if angle > 180.0:
-        angle = 360 - angle
-    return angle
-
-def calculate_lateral_raise_angle(shoulder, wrist):
-    """Calculate angle for lateral raise."""
-    horizontal_reference = np.array([1, 0])
-    arm_vector = np.array([wrist[0] - shoulder[0], wrist[1] - shoulder[1]])
-    
-    dot_product = np.dot(horizontal_reference, arm_vector)
-    magnitude_reference = np.linalg.norm(horizontal_reference)
-    magnitude_arm = np.linalg.norm(arm_vector)
-    
-    if magnitude_arm == 0 or magnitude_reference == 0:
-        return 0
-    
-    cos_angle = dot_product / (magnitude_reference * magnitude_arm)
-    angle = np.arccos(np.clip(cos_angle, -1.0, 1.0))
-    return np.degrees(angle)
-
-class VideoTransformer(VideoTransformerBase):
+class PoseTracker:
     def __init__(self):
-        self.pose = mp_pose.Pose(min_detection_confidence=0.5, min_tracking_confidence=0.5)
-        self.workout_type = "bicep_curl"  # Default workout
+        self.pose = mp_pose.Pose(
+            min_detection_confidence=0.5,
+            min_tracking_confidence=0.5
+        )
+        self.counter = 0
+        self.stage = None
+        self.feedback = ""
+        
+    def calculate_angle(self, a: np.ndarray, b: np.ndarray, c: np.ndarray) -> float:
+        a = np.array(a)
+        b = np.array(b)
+        c = np.array(c)
+        
+        radians = np.arctan2(c[1]-b[1], c[0]-b[0]) - \
+                 np.arctan2(a[1]-b[1], a[0]-b[0])
+        angle = np.abs(np.degrees(radians))
         
-    def process_bicep_curl(self, landmarks):
-        """Process frame for bicep curl exercise."""
+        if angle > 180.0:
+            angle = 360-angle
+            
+        return angle
+
+    def process_frame(self, frame: np.ndarray, exercise_type: str) -> Tuple[np.ndarray, str]:
+        try:
+            # Convert BGR to RGB
+            image = frame.copy()
+            image.flags.writeable = False
+            image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+            results = self.pose.process(image)
+            
+            # Draw landmarks
+            image.flags.writeable = True
+            image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
+            
+            if results.pose_landmarks:
+                mp_drawing.draw_landmarks(
+                    image,
+                    results.pose_landmarks,
+                    mp_pose.POSE_CONNECTIONS
+                )
+                
+                # Extract landmarks
+                try:
+                    landmarks = results.pose_landmarks.landmark
+                    
+                    # Process based on exercise type
+                    if exercise_type == "bicep_curl":
+                        feedback = self._process_bicep_curl(landmarks, image)
+                    elif exercise_type == "shoulder_press":
+                        feedback = self._process_shoulder_press(landmarks, image)
+                    elif exercise_type == "lateral_raise":
+                        feedback = self._process_lateral_raise(landmarks, image)
+                    else:
+                        feedback = "Unknown exercise type"
+                        
+                    return image, feedback
+                    
+                except Exception as e:
+                    logger.error(f"Error processing landmarks: {str(e)}")
+                    return image, "Error processing pose"
+            
+            return image, "No pose detected"
+            
+        except Exception as e:
+            logger.error(f"Error in process_frame: {str(e)}")
+            return frame, "Error processing frame"
+
+    def _process_bicep_curl(self, landmarks, image) -> str:
+        # Get coordinates
         shoulder = [landmarks[mp_pose.PoseLandmark.LEFT_SHOULDER.value].x,
                    landmarks[mp_pose.PoseLandmark.LEFT_SHOULDER.value].y]
         elbow = [landmarks[mp_pose.PoseLandmark.LEFT_ELBOW.value].x,
@@ -101,41 +136,59 @@ class VideoTransformer(VideoTransformerBase):
         wrist = [landmarks[mp_pose.PoseLandmark.LEFT_WRIST.value].x,
                 landmarks[mp_pose.PoseLandmark.LEFT_WRIST.value].y]
         
-        angle = calculate_angle(shoulder, elbow, wrist)
-        
-        with lock:
-            if angle > 160 and state.stage != "down":
-                state.stage = "down"
-                state.feedback = "Lower the weight"
-            elif angle < 40 and state.stage == "down":
-                state.stage = "up"
-                state.counter += 1
-                state.feedback = f"Good rep! Count: {state.counter}"
+        # Calculate angle
+        angle = self.calculate_angle(shoulder, elbow, wrist)
         
-        return angle
+        # Counter logic
+        if angle > 160:
+            self.stage = "down"
+        elif angle < 30 and self.stage == "down":
+            self.stage = "up"
+            self.counter += 1
+            
+        # Add text to image
+        cv2.putText(image, f'Angle: {angle:.2f}', (10,30), 
+                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255,255,255), 2)
+        cv2.putText(image, f'Count: {self.counter}', (10,60),
+                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255,255,255), 2)
+                    
+        if angle < 160 and self.stage == "down":
+            return "Curl up"
+        elif angle > 30 and self.stage == "up":
+            return "Lower the weight"
+        else:
+            return f"Count: {self.counter}"
 
-    def process_lateral_raise(self, landmarks):
-        """Process frame for lateral raise exercise."""
+    def _process_shoulder_press(self, landmarks, image) -> str:
+        # Similar structure to bicep curl but with shoulder press specific angles
         shoulder = [landmarks[mp_pose.PoseLandmark.LEFT_SHOULDER.value].x,
                    landmarks[mp_pose.PoseLandmark.LEFT_SHOULDER.value].y]
+        elbow = [landmarks[mp_pose.PoseLandmark.LEFT_ELBOW.value].x,
+                landmarks[mp_pose.PoseLandmark.LEFT_ELBOW.value].y]
         wrist = [landmarks[mp_pose.PoseLandmark.LEFT_WRIST.value].x,
                 landmarks[mp_pose.PoseLandmark.LEFT_WRIST.value].y]
         
-        angle = calculate_lateral_raise_angle(shoulder, wrist)
+        angle = self.calculate_angle(shoulder, elbow, wrist)
         
-        with lock:
-            if angle < 20 and state.stage != "down":
-                state.stage = "down"
-                state.feedback = "Raise your arms"
-            elif 70 <= angle <= 110 and state.stage == "down":
-                state.stage = "up"
-                state.counter += 1
-                state.feedback = f"Good rep! Count: {state.counter}"
+        if angle < 90:
+            self.stage = "down"
+        elif angle > 160 and self.stage == "down":
+            self.stage = "up"
+            self.counter += 1
+            
+        cv2.putText(image, f'Angle: {angle:.2f}', (10,30),
+                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255,255,255), 2)
+        cv2.putText(image, f'Count: {self.counter}', (10,60),
+                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255,255,255), 2)
         
-        return angle
+        if angle > 90 and self.stage == "down":
+            return "Press up fully"
+        elif angle < 160 and self.stage == "up":
+            return "Lower the weight"
+        else:
+            return f"Count: {self.counter}"
 
-    def process_shoulder_press(self, landmarks):
-        """Process frame for shoulder press exercise."""
+    def _process_lateral_raise(self, landmarks, image) -> str:
         shoulder = [landmarks[mp_pose.PoseLandmark.LEFT_SHOULDER.value].x,
                    landmarks[mp_pose.PoseLandmark.LEFT_SHOULDER.value].y]
         elbow = [landmarks[mp_pose.PoseLandmark.LEFT_ELBOW.value].x,
@@ -143,118 +196,136 @@ class VideoTransformer(VideoTransformerBase):
         wrist = [landmarks[mp_pose.PoseLandmark.LEFT_WRIST.value].x,
                 landmarks[mp_pose.PoseLandmark.LEFT_WRIST.value].y]
         
-        angle = calculate_angle(shoulder, elbow, wrist)
+        angle = self.calculate_angle(shoulder, elbow, wrist)
         
-        with lock:
-            if 80 <= angle <= 100 and state.stage != "down":
-                state.stage = "down"
-                state.feedback = "Press up!"
-            elif angle > 160 and state.stage == "down":
-                state.stage = "up"
-                state.counter += 1
-                state.feedback = f"Good rep! Count: {state.counter}"
+        if angle < 20:
+            self.stage = "down"
+        elif angle > 80 and self.stage == "down":
+            self.stage = "up"
+            self.counter += 1
+            
+        cv2.putText(image, f'Angle: {angle:.2f}', (10,30),
+                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255,255,255), 2)
+        cv2.putText(image, f'Count: {self.counter}', (10,60),
+                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255,255,255), 2)
         
-        return angle
+        if angle < 80 and self.stage == "down":
+            return "Raise arms higher"
+        elif angle > 20 and self.stage == "up":
+            return "Lower arms"
+        else:
+            return f"Count: {self.counter}"
 
-    def recv(self, frame):
-        img = frame.to_ndarray(format="bgr24")
+class VideoProcessor:
+    def __init__(self) -> None:
+        self.pose_tracker = PoseTracker()
+        self._exercise_type = "bicep_curl"
+        
+    @property
+    def exercise_type(self) -> str:
+        return self._exercise_type
         
-        # Process the image
-        image = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
-        results = self.pose.process(image)
+    @exercise_type.setter
+    def exercise_type(self, value: str) -> None:
+        self._exercise_type = value
         
-        if results.pose_landmarks:
-            # Draw pose landmarks
-            mp_drawing.draw_landmarks(
-                img, results.pose_landmarks, mp_pose.POSE_CONNECTIONS,
-                mp_drawing.DrawingSpec(color=(0, 255, 0), thickness=2, circle_radius=2),
-                mp_drawing.DrawingSpec(color=(0, 255, 0), thickness=2)
-            )
+    def recv(self, frame: av.VideoFrame) -> av.VideoFrame:
+        try:
+            img = frame.to_ndarray(format="bgr24")
             
-            # Process based on workout type
-            if self.workout_type == "bicep_curl":
-                angle = self.process_bicep_curl(results.pose_landmarks.landmark)
-            elif self.workout_type == "lateral_raise":
-                angle = self.process_lateral_raise(results.pose_landmarks.landmark)
-            else:  # shoulder_press
-                angle = self.process_shoulder_press(results.pose_landmarks.landmark)
+            # Process the frame
+            processed_frame, feedback = self.pose_tracker.process_frame(img, self.exercise_type)
             
-            # Draw angle and counter
-            cv2.putText(img, f"Angle: {angle:.2f}", (10, 30),
-                       cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2)
-            cv2.putText(img, f"Counter: {state.counter}", (10, 70),
-                       cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2)
-            cv2.putText(img, f"Feedback: {state.feedback}", (10, 110),
-                       cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2)
-        
-        return av.VideoFrame.from_ndarray(img, format="bgr24")
+            # Add feedback to frame
+            cv2.putText(processed_frame, feedback, (10,90),
+                       cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255,255,255), 2)
+            
+            return av.VideoFrame.from_ndarray(processed_frame, format="bgr24")
+        except Exception as e:
+            logger.error(f"Error in recv: {str(e)}")
+            return frame
 
 def main():
+    # Title and description
     st.title("🏋️‍♂️ AI Workout Trainer")
-    
     st.markdown("""
-    <div class='workout-container'>
-        Welcome to your AI Workout Trainer! This app will help you perfect your form
-        and track your exercises in real-time. Choose a workout and follow the feedback
-        to improve your technique.
+    <div class="workout-container">
+        <p>Welcome to your personal AI workout trainer! This app will help you:
+        <ul>
+            <li>Track your exercise form in real-time</li>
+            <li>Count your reps automatically</li>
+            <li>Provide instant feedback on your technique</li>
+        </ul></p>
     </div>
     """, unsafe_allow_html=True)
     
-    # Workout selection
-    workout_options = {
-        "Bicep Curl": "bicep_curl",
-        "Lateral Raise": "lateral_raise",
-        "Shoulder Press": "shoulder_press"
-    }
-    
-    selected_workout = st.selectbox(
-        "Choose your workout:",
-        list(workout_options.keys())
+    # Exercise selection
+    exercise_type = st.selectbox(
+        "Choose your exercise:",
+        ["bicep_curl", "shoulder_press", "lateral_raise"],
+        format_func=lambda x: x.replace('_', ' ').title()
     )
     
-    # Reset state when workout changes
-    if 'last_workout' not in st.session_state or st.session_state.last_workout != selected_workout:
-        with lock:
-            state.counter = 0
-            state.stage = None
-            state.feedback = ""
-        st.session_state.last_workout = selected_workout
-
-    # Exercise descriptions
-    descriptions = {
-        "Bicep Curl": "Focus on keeping your upper arm still and curl the weight up smoothly.",
-        "Lateral Raise": "Raise your arms to shoulder height, keeping them slightly bent.",
-        "Shoulder Press": "Press the weight overhead, fully extending your arms."
-    }
-    
-    st.markdown(f"""
-    <div class='workout-container'>
-        <h3>{selected_workout}</h3>
-        <p>{descriptions[selected_workout]}</p>
-    </div>
-    """, unsafe_allow_html=True)
-
-    # Initialize WebRTC streamer
-    webrtc_ctx = webrtc_streamer(
-        key="workout",
-        video_transformer_factory=VideoTransformer,
-        rtc_configuration={"iceServers": [{"urls": ["stun:stun.l.google.com:19302"]}]}
+    # Create VideoProcessor instance
+    ctx = webrtc_streamer(
+        key="workout-tracker",
+        mode=WebRtcMode.SENDRECV,
+        rtc_configuration=RTC_CONFIGURATION,
+        video_processor_factory=VideoProcessor,
+        media_stream_constraints={"video": True, "audio": False},
+        async_processing=True,
     )
     
-    if webrtc_ctx.video_transformer:
-        webrtc_ctx.video_transformer.workout_type = workout_options[selected_workout]
+    # Update exercise type when changed
+    if ctx.video_processor:
+        ctx.video_processor.exercise_type = exercise_type
     
-    # Display feedback
-    feedback_placeholder = st.empty()
-    if webrtc_ctx.state.playing:
-        feedback_placeholder.markdown(f"""
-        <div class='feedback-text'>
-            <h4>Current Exercise: {selected_workout}</h4>
-            <p>Reps Completed: {state.counter}</p>
-            <p>Feedback: {state.feedback}</p>
+    # Instructions based on exercise
+    if exercise_type == "bicep_curl":
+        st.markdown("""
+        <div class="workout-container">
+            <h3>Bicep Curl Instructions:</h3>
+            <ul>
+                <li>Stand straight with weights at your sides</li>
+                <li>Keep your upper arms still</li>
+                <li>Curl the weights up towards your shoulders</li>
+                <li>Lower the weights back down controlled</li>
+            </ul>
+        </div>
+        """, unsafe_allow_html=True)
+    elif exercise_type == "shoulder_press":
+        st.markdown("""
+        <div class="workout-container">
+            <h3>Shoulder Press Instructions:</h3>
+            <ul>
+                <li>Start with weights at shoulder height</li>
+                <li>Press weights straight up overhead</li>
+                <li>Keep your core tight</li>
+                <li>Lower weights back to shoulders controlled</li>
+            </ul>
+        </div>
+        """, unsafe_allow_html=True)
+    else:  # lateral_raise
+        st.markdown("""
+        <div class="workout-container">
+            <h3>Lateral Raise Instructions:</h3>
+            <ul>
+                <li>Stand straight with weights at your sides</li>
+                <li>Raise arms out to sides up to shoulder height</li>
+                <li>Keep a slight bend in your elbows</li>
+                <li>Lower weights back down controlled</li>
+            </ul>
         </div>
         """, unsafe_allow_html=True)
 
 if __name__ == "__main__":
-    main()
+    try:
+        # Import OpenCV
+        import cv2
+        main()
+    except Exception as e:
+        st.error(f"An error occurred: {str(e)}")
+        if "cv2" in str(e):
+            st.warning("OpenCV import failed. Please check your installation.")
+        logger.error(f"Application error: {str(e)}")