Update app.py

app.py

@@ -1,71 +1,53 @@
 import cv2
 import torch
-import numpy as np
 import gradio as gr
 from ultralytics import YOLO
-import threading
-import time
-
-# Load YOLOv5 model
-device = 'cuda' if torch.cuda.is_available() else 'cpu'
-model = YOLO("yolov5s.pt").to(device)
+import numpy as np
 
-# Open webcam
-cap = cv2.VideoCapture(0)
-frame = np.zeros((480, 640, 3), dtype=np.uint8)  # Default blank frame
-lock = threading.Lock()
+# Load YOLOv5 model (assuming weights are already downloaded)
+model = YOLO("yolov5s.pt")  # You can change to 'yolov5m.pt' or 'yolov5l.pt' for better accuracy
 
-def detect_objects(image):
-    """Detect objects in an uploaded image with YOLO."""
-    image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)  # Convert to BGR for OpenCV
-    results = model.predict(image, conf=0.4)  # Set confidence threshold
-    image = results[0].plot()  # Plot detections directly on image
-    return cv2.cvtColor(image, cv2.COLOR_BGR2RGB)  # Convert back to RGB for Gradio
+def detect_objects_image(image):
+    results = model(image)
+    result_img = results[0].plot()  # Render image with bounding boxes
+    return result_img
 
-def process_webcam():
-    """Continuously capture and process frames from the webcam."""
-    global frame
+# Video detection function
+def detect_objects_video():
+    cap = cv2.VideoCapture(0)  # Capture from default webcam
+    cap.set(cv2.CAP_PROP_FPS, 30)  # Set FPS
+    
     while cap.isOpened():
-        ret, img = cap.read()
+        ret, frame = cap.read()
         if not ret:
-            continue
-
-        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)  # Convert to RGB
-        results = model.predict(img, conf=0.4)  # Explicitly call predict
-        img = results[0].plot()  # Directly draw detections on the frame
+            break
         
-        with lock:
-            frame = img  # Update frame with detection overlay
-
-# Start the webcam processing thread
-threading.Thread(target=process_webcam, daemon=True).start()
+        results = model(frame)
+        result_img = results[0].plot()
+        
+        _, buffer = cv2.imencode(".jpg", result_img)
+        yield buffer.tobytes()
+    
+    cap.release()
 
-def get_webcam_frame():
-    """Returns the latest processed webcam frame."""
-    with lock:
-        return frame
+def start_video():
+    return gr.Video(update=detect_objects_video, streaming=True)
 
 # Gradio UI
 with gr.Blocks() as demo:
-    gr.Markdown("# YOLOv5 Real-Time Object Detection")
-
-    with gr.Tabs():
-        with gr.Tab("Real-Time Webcam"):
-            webcam_output = gr.Image(label="Live Webcam Feed", type="numpy")
-
-            def update_webcam():
-                while True:
-                    with lock:
-                        img = frame
-                    webcam_output.update(img)
-                    time.sleep(1 / 30)  # ~30 FPS
-
-            threading.Thread(target=update_webcam, daemon=True).start()
-
-        with gr.Tab("Upload Image"):
-            image_input = gr.Image(type="numpy", label="Upload Image")
-            image_output = gr.Image(label="Detected Objects")
-            image_button = gr.Button("Detect Objects")
-            image_button.click(detect_objects, inputs=image_input, outputs=image_output)
+    gr.Markdown("## Live Object Detection with YOLOv5")
+    
+    with gr.Row():
+        img_input = gr.Image(type="numpy")
+        img_output = gr.Image()
+        img_button = gr.Button("Detect Objects in Image")
+    
+    img_button.click(detect_objects_image, inputs=img_input, outputs=img_output)
+    
+    with gr.Row():
+        video_button = gr.Button("Start Live Video Detection")
+        video_output = gr.Video()
+    
+    video_button.click(start_video, outputs=video_output)
 
 demo.launch()