Update app.py

app.py

@@ -25,28 +25,36 @@ transform = transforms.Compose([
 
 # Prediction function
 def predict(image):
-    image = transform(image).unsqueeze(0)  # Add batch dimension
-    with torch.no_grad():
-        predictions = model(image)
-    
-    # Process predictions
-    boxes = predictions[0]['boxes'].cpu().numpy()
-    scores = predictions[0]['scores'].cpu().numpy()
-    labels = predictions[0]['labels'].cpu().numpy()
-
-    # Filter out low-confidence predictions
-    threshold = 0.5
-    boxes = boxes[scores > threshold]
-    labels = labels[scores > threshold]
-
-    # Draw boxes on the image
-    image_np = np.array(image.squeeze().permute(1, 2, 0).cpu())
-    for box, label in zip(boxes, labels):
-        x1, y1, x2, y2 = box.astype(int)
-        image_np = cv2.rectangle(image_np, (x1, y1), (x2, y2), (255, 0, 0), 2)
-        image_np = cv2.putText(image_np, str(label), (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 2)
-
-    return Image.fromarray(image_np)
+    try:
+        # Transform the image
+        image_tensor = transform(image).unsqueeze(0)  # Add batch dimension
+        with torch.no_grad():
+            predictions = model(image_tensor)
+
+        # Process predictions
+        boxes = predictions[0]['boxes'].cpu().numpy()
+        scores = predictions[0]['scores'].cpu().numpy()
+        labels = predictions[0]['labels'].cpu().numpy()
+
+        # Filter out low-confidence predictions
+        threshold = 0.5
+        boxes = boxes[scores > threshold]
+        labels = labels[scores > threshold]
+
+        # Convert the input image to a NumPy array
+        image_np = np.array(image)
+
+        # Draw boxes on the image
+        for box, label in zip(boxes, labels):
+            x1, y1, x2, y2 = box.astype(int)
+            image_np = cv2.rectangle(image_np, (x1, y1), (x2, y2), (255, 0, 0), 2)
+            image_np = cv2.putText(image_np, str(label), (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 2)
+
+        # Ensure the output is in the correct format
+        return Image.fromarray(image_np.astype(np.uint8))
+
+    except Exception as e:
+        return f"Error: {str(e)}"
 
 # Gradio interface
 iface = gr.Interface(fn=predict, inputs=gr.Image(type="pil"), outputs=gr.Image(type="pil"), title="Object Detection with Faster R-CNN")