Update app.py

app.py

@@ -5,8 +5,6 @@ import torchvision.transforms as transforms
 import base64
 import cv2
 import numpy as np
-from streamlit_webrtc import webrtc_streamer
-import av
 
 # Set Streamlit Page Configuration
 st.set_page_config(
@@ -47,12 +45,32 @@ def get_base64_image(image_path):
     except FileNotFoundError:
         return None
 
-# Function for real-time PPE detection using WebRTC
-def process_frame(frame):
-    img = frame.to_ndarray(format="bgr24")  # Convert frame to numpy array
-    results = model.predict(img)  # Run detection
-    output_img = results[0].plot()  # Draw detection results
-    return av.VideoFrame.from_ndarray(output_img, format="bgr24")  # Convert back to VideoFrame
+# Function for real-time PPE detection using webcam
+def live_ppe_detection():
+    st.sidebar.write("Starting live detection...")
+    cap = cv2.VideoCapture(0)
+    if not cap.isOpened():
+        st.sidebar.error("Error: Could not open webcam.")
+        return
+    
+    stframe = st.empty()
+    stop_button = st.sidebar.button("Stop Live Detection", key="stop_button")
+    
+    while cap.isOpened():
+        ret, frame = cap.read()
+        if not ret:
+            st.sidebar.error("Failed to capture video frame.")
+            break
+        
+        results = model.predict(frame)
+        output_frame = results[0].plot()
+        stframe.image(output_frame, channels="BGR")
+        
+        if stop_button:
+            break
+    
+    cap.release()
+    cv2.destroyAllWindows()
 
 # Display logo
 image_base64 = get_base64_image("logo/logo.png")
@@ -79,8 +97,8 @@ uploaded_file = st.sidebar.file_uploader("Drag and drop or browse", type=['jpg',
 
 # Sidebar - Live Predictions
 st.sidebar.header("📡 Live Predictions")
-st.sidebar.write("Start real-time PPE detection using your webcam")
-webrtc_streamer(key="live", video_frame_callback=process_frame)
+if st.sidebar.button("Start Live Detection", key="start_button"):
+    live_ppe_detection()
 
 # Main Page
 st.title("PPE Detect")
@@ -102,6 +120,4 @@ if uploaded_file:
             st.error("Detection failed. Please try again.")
 
 
-
-
 st.info("This app uses **YOLO** for PPE detection. Upload an image or start live detection to get started.")