Update app.py

app.py

@@ -1,57 +1,47 @@
 import streamlit as st
 import torch
-from PIL import Image
-from torchvision.transforms import functional as F
-from yolov5.utils.general import non_max_suppression, scale_coords
-from yolov5.models.experimental import attempt_load
-from yolov5.utils.plots import plot_one_box
 import cv2
+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
+
+# Streamlit UI
+st.title('YOLOv5 Object Detection')
+
+# Input Options
+upload_option = st.radio("Choose an input option:", ("Upload Image", "Real-Time Webcam"))
+
+# Image Upload Handling
+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
+
+# Real-Time Webcam Detection
+if upload_option == "Real-Time Webcam":
+    run = st.checkbox('Run Webcam')
+    FRAME_WINDOW = st.image([])  # Display window for webcam frames
+
+    if run:
+        cap = cv2.VideoCapture(0)  # Open webcam (0 for default)
+
+        while run:
+            ret, frame = cap.read()  # Capture frame
+            if not ret:
+                st.write("Error: Unable to capture frame")
+                break
+                
+            # Convert to RGB and detect objects
+            frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+            results = model(frame_rgb)
+
+            annotated_frame = results.render()[0]
+            FRAME_WINDOW.image(annotated_frame)
+
+        cap.release()  # Release webcam
 
-@st.cache(allow_output_mutation=True)
-def load_model():
-    # Load your pre-trained YOLOv5 model
-    model = attempt_load('best.pt', map_location=torch.device('cpu'))
-    return model
-
-def detect_objects(image, model, confidence=0.4, iou=0.5):
-    img = Image.fromarray(image.astype('uint8')).convert('RGB')
-    img_tensor = F.to_tensor(img)
-    img_tensor, _ = model.preprocess(img_tensor, None, None)
-    pred = model(img_tensor)[0]
-    pred = non_max_suppression(pred, confidence, iou)[0]
-
-    if pred is not None and len(pred):
-        pred[:, :4] = scale_coords(img_tensor.shape[2:], pred[:, :4], img.size).round()
-
-    return pred
-
-def main():
-    st.title("Real-time Object Detection with YOLOv5")
-
-    # Choose between image upload or video stream
-    option = st.radio("Choose Input Type:", ("Image Upload", "Video Stream"))
-
-    if option == "Image Upload":
-        uploaded_image = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
-
-        if uploaded_image is not None:
-            image = Image.open(uploaded_image)
-            st.image(image, caption="Uploaded Image", use_column_width=True)
-            st.write("")
-
-            if st.button("Detect Objects"):
-                st.write("Detecting...")
-                model = load_model()
-                with st.spinner('Wait for it...'):
-                    pred = detect_objects(image, model)
-                    if pred is not None and len(pred):
-                        for *xyxy, conf, cls in pred:
-                            label = f'{model.names[int(cls)]} {conf:.2f}'
-                            plot_one_box(xyxy, image, label=label, color=(255, 0, 0), line_thickness=2)
-                        st.image(image, caption="Result", use_column_width=True)
-
-    elif option == "Video Stream":
-        st.write("Video stream functionality is not implemented yet.")
-
-if __name__ == '__main__':
-    main()