Output Return an image

app.py

@@ -1,6 +1,6 @@
 import gradio as gr
 from huggingface_hub import hf_hub_download
-from PIL import Image
+from PIL import Image, ImageDraw
 import torch
 from transformers import AutoImageProcessor, AutoModelForObjectDetection
 
@@ -20,16 +20,27 @@ def predict(image):
     # Extract bounding boxes and class labels
     predicted_boxes = outputs.pred_boxes[0].cpu().numpy()  # First image
     predicted_classes = outputs.logits.argmax(-1).cpu().numpy()  # Class predictions
-
-    # Return the bounding boxes for display
-    return {"boxes": predicted_boxes.tolist(), "classes": predicted_classes.tolist()}
+    
+    # Create a drawing context for the image
+    draw = ImageDraw.Draw(image)
+    width, height = image.size
+    
+    # Loop over all detected boxes and draw them on the image
+    for box in predicted_boxes:
+        # Box coordinates are normalized, so multiply by image dimensions
+        x0, y0, x1, y1 = box
+        draw.rectangle([x0 * width, y0 * height, x1 * width, y1 * height], outline="red", width=3)
+    
+    # Return the image with bounding boxes drawn
+    return image
 
 # Set up the Gradio interface
 interface = gr.Interface(
     fn=predict,  # The function that gets called when an image is uploaded
     inputs=gr.Image(type="pil"),  # Image input (as PIL image)
-    outputs="json",  # Outputting a JSON with the boxes and classes
+    outputs=gr.Image(type="pil"),  # Outputting the image with boxes drawn
 )
 
 # Launch the Gradio app
 interface.launch()
+