Update app.py

app.py

@@ -21,33 +21,47 @@ model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True).to(devi
 model.eval()
 
 def preprocess_image(image):
-    image_tensor = F.to_tensor(image)
-    return image_tensor.unsqueeze(0).to(device)
+    try:
+        image_tensor = F.to_tensor(image)
+        return image_tensor.unsqueeze(0).to(device)
+    except Exception as e:
+        print(f"Error in preprocessing image: {e}")
+        return None
 
 def draw_boxes(image, outputs, threshold=0.3):
-    image = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
-    h, w, _ = image.shape
+    try:
+        image = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
+        h, w, _ = image.shape
 
-    for box in outputs:
-        score, label, x1, y1, x2, y2 = box[4].item(), int(box[5].item()), box[0].item(), box[1].item(), box[2].item(), box[3].item()
-        if score > threshold:
-            x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
-            cv2.rectangle(image, (x1, y1), (x2, y2), (255, 0, 0), 2)
-            text = f"{model.names[label]:s}: {score:.2f}"
-            cv2.putText(image, text, (x1, y1 - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 2)
+        for box in outputs:
+            score, label, x1, y1, x2, y2 = box[4].item(), int(box[5].item()), box[0].item(), box[1].item(), box[2].item(), box[3].item()
+            if score > threshold:
+                x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
+                cv2.rectangle(image, (x1, y1), (x2, y2), (255, 0, 0), 2)
+                text = f"{model.names[label]:s}: {score:.2f}"
+                cv2.putText(image, text, (x1, y1 - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 2)
 
-    return cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+        return cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+    except Exception as e:
+        print(f"Error in drawing boxes: {e}")
+        return image
 
 def detect_objects(image):
     image_tensor = preprocess_image(image)
-    outputs = model(image_tensor)
-    outputs = non_max_suppression(outputs)[0]
-    result_image = draw_boxes(image, outputs)
-    return result_image
+    if image_tensor is None:
+        return image
+    try:
+        outputs = model(image_tensor)
+        outputs = non_max_suppression(outputs)[0]
+        result_image = draw_boxes(image, outputs)
+        return result_image
+    except Exception as e:
+        print(f"Error in detecting objects: {e}")
+        return image
 
 iface = gr.Interface(
     fn=detect_objects,
-    inputs=gr.Image(type="pil"),
+    inputs=gr.Video(source="webcam", type="pil"),
     outputs=gr.Image(type="pil"),
     live=True
 )