Update app.py

app.py

@@ -5,8 +5,13 @@ from PIL import Image
 import numpy as np
 import tempfile
 
+
 # YOLOv5 Model Loading (best.pt)
-model = torch.hub.load('ultralytics/yolov5', 'custom', path='best.pt')  # Replace 'best.pt' with your model's path
+@st.cache_resource  # Cache the model for efficiency
+def load_model():
+    return torch.hub.load('ultralytics/yolov5', 'custom', path='best.pt', force_reload=True) 
+
+model = load_model()
 
 # Streamlit UI
 st.title('YOLOv5 Object Detection')
@@ -19,29 +24,33 @@ if upload_option == "Upload Image":
     uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "png", "jpeg"])
     if uploaded_file is not None:
         image = Image.open(uploaded_file)
-        results = model(image)  # Perform inference
-        st.image(results.render()[0], caption='Detected Objects', use_column_width=True)  # Display results
+        results = model(image)  
+        st.image(results.render()[0], caption='Detected Objects', use_column_width=True) 
 
 # Real-Time Webcam Detection
 if upload_option == "Real-Time Webcam":
     run = st.checkbox('Run Webcam')
-    FRAME_WINDOW = st.image([])  # Display window for webcam frames
+    FRAME_WINDOW = st.image([])
 
     if run:
-        cap = cv2.VideoCapture(0)  # Open webcam (0 for default)
+        try:
+            cap = cv2.VideoCapture(0) 
+        except Exception as e:  
+            st.error(f"Error accessing webcam: {e}")
+            st.stop()
 
         while run:
-            ret, frame = cap.read()  # Capture frame
+            ret, frame = cap.read()
             if not ret:
-                st.write("Error: Unable to capture frame")
-                break
-                
+                st.warning("Error: Unable to capture frame. Please check your webcam settings.")  # More informative message
+                continue  # Skip to the next iteration instead of stopping the loop entirely
+
             # Convert to RGB and detect objects
             frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
             results = model(frame_rgb)
 
+            # Render and display results
             annotated_frame = results.render()[0]
-            FRAME_WINDOW.image(annotated_frame)
-
-        cap.release()  # Release webcam
+            FRAME_WINDOW.image(annotated_frame, channels="BGR")  # Display in BGR format
 
+        cap.release()