Update app.py

app.py

@@ -10,32 +10,18 @@ import time
 device = 'cuda' if torch.cuda.is_available() else 'cpu'
 model = YOLO("yolov5s.pt").to(device)
 
-# Generate unique colors for each class
-num_classes = len(model.names)
-colors = np.random.randint(0, 255, size=(num_classes, 3), dtype=np.uint8)
+# Open webcam
+cap = cv2.VideoCapture(0)
+frame = np.zeros((480, 640, 3), dtype=np.uint8)  # Default blank frame
+lock = threading.Lock()
 
 def detect_objects(image):
-    """Detect objects in an uploaded image with different bounding box colors."""
+    """Detect objects in an uploaded image with YOLO."""
     image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)  # Convert to BGR for OpenCV
-    results = model.predict(image)  # Explicitly call predict
-    detections = results[0].boxes.data.cpu().numpy()  # Get detections
-    
-    for box in detections:
-        x1, y1, x2, y2, conf, cls = box[:6]
-        cls = int(cls)  # Convert class to int
-        label = f"{model.names[cls]} {conf:.2f}"
-        color = tuple(map(int, colors[cls]))  # Assign unique color
-        
-        cv2.rectangle(image, (int(x1), int(y1)), (int(x2), int(y2)), color, 2)
-        cv2.putText(image, label, (int(x1), int(y1) - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
-    
+    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
 
-# Real-time webcam processing
-cap = cv2.VideoCapture(0)
-frame = None
-lock = threading.Lock()
-
 def process_webcam():
     """Continuously capture and process frames from the webcam."""
     global frame
@@ -43,30 +29,21 @@ def process_webcam():
         ret, img = cap.read()
         if not ret:
             continue
-        
+
         img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)  # Convert to RGB
-        results = model.predict(img)  # Explicitly call predict
-        detections = results[0].boxes.data.cpu().numpy()
-        
-        for box in detections:
-            x1, y1, x2, y2, conf, cls = box[:6]
-            cls = int(cls)  # Convert class to int
-            label = f"{model.names[cls]} {conf:.2f}"
-            color = tuple(map(int, colors[cls]))  # Assign unique color
-        
-            cv2.rectangle(img, (int(x1), int(y1)), (int(x2), int(y2)), color, 2)
-            cv2.putText(img, label, (int(x1), int(y1) - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
+        results = model.predict(img, conf=0.4)  # Explicitly call predict
+        img = results[0].plot()  # Directly draw detections on the frame
         
         with lock:
-            frame = img
+            frame = img  # Update frame with detection overlay
 
-# Start the webcam thread
+# Start the webcam processing thread
 threading.Thread(target=process_webcam, daemon=True).start()
 
 def get_webcam_frame():
     """Returns the latest processed webcam frame."""
     with lock:
-        return frame if frame is not None else np.zeros((480, 640, 3), dtype=np.uint8)
+        return frame
 
 # Gradio UI
 with gr.Blocks() as demo:
@@ -74,14 +51,14 @@ with gr.Blocks() as demo:
 
     with gr.Tabs():
         with gr.Tab("Real-Time Webcam"):
-            webcam_output = gr.Image(label="Live Webcam Feed")
-            
+            webcam_output = gr.Image(label="Live Webcam Feed", type="numpy")
+
             def update_webcam():
                 while True:
                     with lock:
-                        img = frame if frame is not None else np.zeros((480, 640, 3), dtype=np.uint8)
+                        img = frame
                     webcam_output.update(img)
-                    time.sleep(1/30)  # ~30 FPS
+                    time.sleep(1 / 30)  # ~30 FPS
 
             threading.Thread(target=update_webcam, daemon=True).start()