Update app.py

app.py

@@ -6,7 +6,7 @@ from ultralytics import YOLO
 import threading
 import time
 
-# Load YOLOv5 model (optimized for CUDA if available)
+# Load YOLOv5 model
 device = 'cuda' if torch.cuda.is_available() else 'cpu'
 model = YOLO("yolov5s.pt").to(device)
 
@@ -16,42 +16,46 @@ colors = np.random.randint(0, 255, size=(num_classes, 3), dtype=np.uint8)
 
 def detect_objects(image):
     """Detect objects in an uploaded image with different bounding box colors."""
-    results = model(image)
+    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 = map(int, box[:6])
+        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 based on class
+        color = tuple(map(int, colors[cls]))  # Assign unique color
         
-        cv2.rectangle(image, (x1, y1), (x2, y2), color, 2)
-        cv2.putText(image, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
+        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)
     
-    return image
+    return cv2.cvtColor(image, cv2.COLOR_BGR2RGB)  # Convert back to RGB for Gradio
 
 # Real-time webcam processing
-cap = cv2.VideoCapture(0)  # Capture from webcam
+cap = cv2.VideoCapture(0)
 frame = None
 lock = threading.Lock()
 
 def process_webcam():
     """Continuously capture and process frames from the webcam."""
     global frame
-    while True:
+    while cap.isOpened():
         ret, img = cap.read()
         if not ret:
             continue
         
-        results = model(img)
+        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 = map(int, box[:6])
+            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, (x1, y1), (x2, y2), color, 2)
-            cv2.putText(img, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
+            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)
         
         with lock:
             frame = img
@@ -71,9 +75,12 @@ with gr.Blocks() as demo:
     with gr.Tabs():
         with gr.Tab("Real-Time Webcam"):
             webcam_output = gr.Image(label="Live Webcam Feed")
+            
             def update_webcam():
                 while True:
-                    webcam_output.update(get_webcam_frame())
+                    with lock:
+                        img = frame if frame is not None else np.zeros((480, 640, 3), dtype=np.uint8)
+                    webcam_output.update(img)
                     time.sleep(1/30)  # ~30 FPS
 
             threading.Thread(target=update_webcam, daemon=True).start()