Update app.py

app.py

@@ -4,21 +4,28 @@ import numpy as np
 import gradio as gr
 from ultralytics import YOLO
 import threading
+import time
 
 # Load YOLOv5 model (optimized for CUDA if available)
 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)
+
 def detect_objects(image):
-    """Detect objects in an uploaded image."""
+    """Detect objects in an uploaded image with different bounding box colors."""
     results = model(image)
     detections = results[0].boxes.data.cpu().numpy()  # Get detections
     
     for box in detections:
         x1, y1, x2, y2, conf, cls = map(int, box[:6])
         label = f"{model.names[cls]} {conf:.2f}"
-        cv2.rectangle(image, (x1, y1), (x2, y2), (0, 255, 0), 2)
-        cv2.putText(image, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
+        color = tuple(map(int, colors[cls]))  # Assign unique color based on class
+        
+        cv2.rectangle(image, (x1, y1), (x2, y2), color, 2)
+        cv2.putText(image, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
     
     return image
 
@@ -28,6 +35,7 @@ frame = None
 lock = threading.Lock()
 
 def process_webcam():
+    """Continuously capture and process frames from the webcam."""
     global frame
     while True:
         ret, img = cap.read()
@@ -40,8 +48,10 @@ def process_webcam():
         for box in detections:
             x1, y1, x2, y2, conf, cls = map(int, box[:6])
             label = f"{model.names[cls]} {conf:.2f}"
-            cv2.rectangle(img, (x1, y1), (x2, y2), (0, 255, 0), 2)
-            cv2.putText(img, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
+            color = tuple(map(int, colors[cls]))  # Assign unique color
+        
+            cv2.rectangle(img, (x1, y1), (x2, y2), color, 2)
+            cv2.putText(img, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
         
         with lock:
             frame = img
@@ -55,16 +65,22 @@ def get_webcam_frame():
         return frame if frame is not None else np.zeros((480, 640, 3), dtype=np.uint8)
 
 # Gradio UI
-demo = gr.Blocks()
-
-with demo:
+with gr.Blocks() as demo:
     gr.Markdown("# YOLOv5 Real-Time Object Detection")
+
     with gr.Tabs():
         with gr.Tab("Real-Time Webcam"):
-            gr.Video(get_webcam_frame, streaming=True)
+            webcam_output = gr.Image(label="Live Webcam Feed")
+            def update_webcam():
+                while True:
+                    webcam_output.update(get_webcam_frame())
+                    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")
-            image_output = gr.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)