Update app.py

app.py

@@ -187,42 +187,45 @@
 #     """
 # ) 
 
+)
+
+import logging
 import cv2
-import mediapipe as mp
 import numpy as np
+import mediapipe as mp
 import streamlit as st
 from streamlit_webrtc import webrtc_streamer
 
-# Initialize MediaPipe Pose
+# Logging setup
+logger = logging.getLogger(__name__)
+
+# Streamlit setup
+st.title("AI Squat Detection using WebRTC")
+st.info("Use your webcam for real-time squat detection.")
+
+# Initialize MediaPipe components
 mp_pose = mp.solutions.pose
 mp_drawing = mp.solutions.drawing_utils
 
-# Function to calculate angles between three points
+# Angle calculation function
 def calculate_angle(a, b, c):
-    a = np.array(a)  # First point
-    b = np.array(b)  # Midpoint
-    c = np.array(c)  # Endpoint
+    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(radians * 180.0 / np.pi)
     if angle > 180.0:
         angle = 360 - angle
     return angle
 
-# Squat detection processor class
-class VideoProcessor:
-    def __init__(self):
-        self.pose = mp_pose.Pose(min_detection_confidence=0.5, min_tracking_confidence=0.5)
-
-    def recv(self, frame):
-        image = frame.to_ndarray(format="bgr24")
-        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
-        image.flags.writeable = False
-        results = self.pose.process(image)
-        image.flags.writeable = True
-        image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
-
-        try:
-            landmarks = results.pose_landmarks.landmark
+def process_frame(frame):
+    image = frame.to_ndarray(format="bgr24")
+    image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+    with mp_pose.Pose(min_detection_confidence=0.5, min_tracking_confidence=0.5) as pose:
+        results = pose.process(image_rgb)
+        landmarks = results.pose_landmarks.landmark if results.pose_landmarks else []
+        
+        if landmarks:
             hip = [landmarks[mp_pose.PoseLandmark.LEFT_HIP.value].x, 
                    landmarks[mp_pose.PoseLandmark.LEFT_HIP.value].y]
             knee = [landmarks[mp_pose.PoseLandmark.LEFT_KNEE.value].x, 
@@ -231,39 +234,43 @@ class VideoProcessor:
                      landmarks[mp_pose.PoseLandmark.LEFT_ANKLE.value].y]
             shoulder = [landmarks[mp_pose.PoseLandmark.LEFT_SHOULDER.value].x, 
                         landmarks[mp_pose.PoseLandmark.LEFT_SHOULDER.value].y]
-            foot_index = [landmarks[mp_pose.PoseLandmark.LEFT_FOOT_INDEX.value].x, 
-                          landmarks[mp_pose.PoseLandmark.LEFT_FOOT_INDEX.value].y]
-            x_axis_hip = [hip[0], 0]
-
-            angle_knee = calculate_angle(hip, knee, ankle)
-            angle_ankle = calculate_angle(foot_index, ankle, knee)
-            angle_hip = calculate_angle(shoulder, hip, x_axis_hip)
-
-            feedback = ""
-            if 80 < angle_knee < 110 and 29 < angle_hip < 40:
-                feedback = "Good Squat!"
-            elif angle_knee < 80:
-                feedback = "Squat too deep!"
-            elif angle_knee > 110:
-                feedback = "Lower your hips!"
-            elif angle_hip < 29:
-                feedback = "Bend Forward!"
-            elif angle_hip > 45:
-                feedback = "Bend Backward!"
-
-            cv2.putText(image, feedback, (50, 50), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2, cv2.LINE_AA)
-            mp_drawing.draw_landmarks(image, results.pose_landmarks, mp_pose.POSE_CONNECTIONS)
-        except Exception as e:
-            print(f"Error: {e}")
-
-        return frame.from_ndarray(image, format="bgr24")
-
-# Streamlit WebRTC configuration
+            foot = [landmarks[mp_pose.PoseLandmark.LEFT_FOOT_INDEX.value].x, 
+                    landmarks[mp_pose.PoseLandmark.LEFT_FOOT_INDEX.value].y]
+
+            # Calculate angles
+            knee_angle = calculate_angle(hip, knee, ankle)
+            hip_angle = calculate_angle(shoulder, hip, [hip[0], 0])
+            ankle_angle = calculate_angle(foot, ankle, knee)
+            
+            # Display key angles
+            cv2.putText(image, f"Knee: {int(knee_angle)}", (10, 50), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2)
+            cv2.putText(image, f"Hip: {int(hip_angle)}", (10, 100), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2)
+            cv2.putText(image, f"Ankle: {int(ankle_angle)}", (10, 150), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2)
+            
+            # Squat logic
+            if 80 < knee_angle < 110 and 29 < hip_angle < 40:
+                cv2.putText(image, "Squat Detected!", (300, 100), cv2.FONT_HERSHEY_SIMPLEX, 2, (0, 255, 0), 3)
+            else:
+                if hip_angle < 29:
+                    cv2.putText(image, "Lean Forward!", (300, 200), cv2.FONT_HERSHEY_SIMPLEX, 2, (0, 0, 255), 3)
+                elif hip_angle > 45:
+                    cv2.putText(image, "Lean Backward!", (300, 200), cv2.FONT_HERSHEY_SIMPLEX, 2, (0, 0, 255), 3)
+                if knee_angle < 80:
+                    cv2.putText(image, "Squat Too Deep!", (300, 250), cv2.FONT_HERSHEY_SIMPLEX, 2, (0, 0, 255), 3)
+                elif knee_angle > 110:
+                    cv2.putText(image, "Lower Your Hips!", (300, 300), cv2.FONT_HERSHEY_SIMPLEX, 2, (0, 0, 255), 3)
+        
+        mp_drawing.draw_landmarks(image, results.pose_landmarks, mp_pose.POSE_CONNECTIONS,
+                                  mp_drawing.DrawingSpec(color=(255, 175, 0), thickness=2, circle_radius=2),
+                                  mp_drawing.DrawingSpec(color=(0, 255, 200), thickness=2, circle_radius=2))
+    return image
+
+# WebRTC streamer configuration
 webrtc_streamer(
-    key="squat_detector",
-    video_processor_factory=VideoProcessor,
+    key="squat-detection",
+    video_processor_factory=lambda: process_frame,
     media_stream_constraints={"video": True, "audio": False},
-    rtc_configuration={"iceServers": [{"urls": ["stun:stun.l.google.com:19302"]}]},
+    async_processing=True
 )
 
 
@@ -292,8 +299,6 @@ webrtc_streamer(
 
 
 
-
-