Update app.py

app.py

@@ -1,6 +1,6 @@
 import gradio as gr
 from transformers import AutoImageProcessor, AutoModelForObjectDetection
-from PIL import Image
+from PIL import Image, ImageDraw
 import torch
 
 image_processor = AutoImageProcessor.from_pretrained('hustvl/yolos-small')
@@ -15,11 +15,16 @@ def detect(image):
     results = image_processor.post_process_object_detection(outputs,
                                                             threshold=0.9,
                                                             target_sizes=target_sizes)[0]
+
+    # Bounding box in COCO format:
+    # [x_min, y_min, width, height]
     
     # model predicts bounding boxes and corresponding COCO classes
     #logits = outputs.logits
     #bboxes = outputs.pred_boxes
 
+    draw = ImageDraw.Draw(image)
+    
     # label and the count
     counts = {}
 
@@ -30,12 +35,16 @@ def detect(image):
             counts[label_name] = 0
         counts[label_name] += 1
 
-    return results
+        x, y, w, h = tuple(box)
+        draw.rectangle((x, y, x+w, y+h), outline="red", width=1)
+        draw.text((x, y), label_name, fill="white")
+        
+    return results, image
 
 demo = gr.Interface(
     fn=detect,
     inputs=[gr.inputs.Image(label="Input image", type="pil")],
-    outputs=["text"], #, gr.Label(num_top_classes=10)],
+    outputs=["text", "image"], #, gr.Label(num_top_classes=10)],
     title="Object Counts in Image"
 )