Update app.py

app.py

@@ -1,8 +1,9 @@
 import streamlit as st
-import cv2
+import av
 import joblib
 import mediapipe as mp
 import numpy as np
+from streamlit_webrtc import webrtc_streamer, VideoTransformerBase
 
 # Load trained model and label encoder
 model = joblib.load("pose_classifier.joblib")
@@ -13,48 +14,37 @@ mp_pose = mp.solutions.pose
 pose = mp_pose.Pose()
 
 # Streamlit UI
-st.title("Live Pose Classification")
-st.write("Real-time pose detection using OpenCV and MediaPipe.")
+st.title("Live Pose Classification on Hugging Face Spaces")
+st.write("Using Streamlit WebRTC, OpenCV, and MediaPipe.")
 
-# OpenCV Video Capture
-cap = cv2.VideoCapture(0)
+class PoseClassification(VideoTransformerBase):
+    def transform(self, frame):
+        img = frame.to_ndarray(format="bgr24")
 
-# Streamlit Image Display
-frame_placeholder = st.empty()
+        # Convert frame to RGB
+        img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
 
-while cap.isOpened():
-    ret, frame = cap.read()
-    if not ret:
-        st.warning("Failed to capture video. Check your camera.")
-        break
+        # Process frame with MediaPipe Pose
+        results = pose.process(img_rgb)
 
-    # Convert frame to RGB
-    img_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+        if results.pose_landmarks:
+            landmarks = results.pose_landmarks.landmark
+            pose_data = [j.x for j in landmarks] + [j.y for j in landmarks] + \
+                        [j.z for j in landmarks] + [j.visibility for j in landmarks]
 
-    # Process frame with MediaPipe Pose
-    results = pose.process(img_rgb)
+            pose_data = np.array(pose_data).reshape(1, -1)
 
-    if results.pose_landmarks:
-        landmarks = results.pose_landmarks.landmark
-        pose_data = [j.x for j in landmarks] + [j.y for j in landmarks] + \
-                    [j.z for j in landmarks] + [j.visibility for j in landmarks]
+            try:
+                y_pred = model.predict(pose_data)
+                predicted_label = label_encoder.inverse_transform(y_pred)[0]
 
-        pose_data = np.array(pose_data).reshape(1, -1)
+                # Draw label on frame
+                cv2.putText(img, f"Pose: {predicted_label}", (20, 50),
+                            cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 3)
+            except Exception as e:
+                st.warning(f"⚠️ Prediction Error: {e}")
 
-        # Predict pose
-        y_pred = model.predict(pose_data)
-        predicted_label = label_encoder.inverse_transform(y_pred)[0]
+        return av.VideoFrame.from_ndarray(img, format="bgr24")
 
-        # Display predicted label
-        cv2.putText(frame, f"Pose: {predicted_label}", (20, 50),
-                    cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 3)
-
-    # Display frame in Streamlit
-    frame_placeholder.image(frame, channels="BGR")
-
-    # Break loop if user stops execution
-    if st.button("Stop Camera"):
-        break
-
-cap.release()
-cv2.destroyAllWindows()
+# Start WebRTC streamer
+webrtc_streamer(key="pose-classification", video_transformer_factory=PoseClassification)