Update app.py

app.py

@@ -5,11 +5,7 @@ from PIL import Image
 import numpy as np
 
 # YOLOv5 Model Loading (best.pt)
-@st.cache(allow_output_mutation=True)  # Cache the model for efficiency
-def load_model():
-    return torch.hub.load('ultralytics/yolov5', 'custom', path='best.pt', force_reload=True) 
-
-model = load_model()
+model = torch.hub.load('ultralytics/yolov5', 'custom', path='best.pt')  # Replace 'best.pt' with your model's path
 
 # Streamlit UI
 st.title('YOLOv5 Object Detection')
@@ -22,39 +18,32 @@ 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)
-        # Convert PIL image to numpy array
-        image_np = np.array(image)
-        # Convert numpy array to PyTorch tensor
-        image_tensor = torch.from_numpy(image_np)
-        results = model(image_tensor)  
-        st.image(results.render()[0], caption='Detected Objects', use_column_width=True) 
+        results = model(image)  # Perform inference
+        st.image(results.render()[0], caption='Detected Objects', use_column_width=True)  # Display results
 
 # Real-Time Webcam Detection
 if upload_option == "Real-Time Webcam":
     run = st.checkbox('Run Webcam')
-    FRAME_WINDOW = st.image([])
+    FRAME_WINDOW = st.image([])  # Display window for webcam frames
 
     if run:
-        try:
-            cap = cv2.VideoCapture(0) 
-        except Exception as e:  
-            st.error(f"Error accessing webcam: {e}")
+        cap = cv2.VideoCapture(0)  # Open webcam (0 for default)
+
+        if not cap.isOpened():
+            st.write("Error: Unable to open webcam")
             st.stop()
 
         while run:
-            ret, frame = cap.read()
+            ret, frame = cap.read()  # Capture frame
             if not ret:
-                st.warning("Error: Unable to capture frame. Please check your webcam settings.")  # More informative message
-                break  # Exit the loop if frame capture fails
+                st.write("Error: Unable to capture frame")
+                break
 
             # Convert to RGB and detect objects
             frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
-            # Convert numpy array to PyTorch tensor
-            frame_tensor = torch.from_numpy(frame_rgb)
-            results = model(frame_tensor)
+            results = model(frame_rgb)
 
-            # Render and display results
             annotated_frame = results.render()[0]
-            FRAME_WINDOW.image(annotated_frame, channels="RGB")  # Display in RGB format
+            FRAME_WINDOW.image(annotated_frame)
 
-        cap.release()
+        cap.release()  # Release webcam