Create app.py

app.py

@@ -0,0 +1,82 @@
+import gradio as gr
+import pandas as pd
+from transformers import AutoImageProcessor, AutoModelForObjectDetection
+from PIL import Image, ImageDraw
+import torch
+
+image_processor = AutoImageProcessor.from_pretrained('hustvl/yolos-small')
+model = AutoModelForObjectDetection.from_pretrained('hustvl/yolos-small')
+
+colors = ["red",
+          "orange",
+          "yellow",
+          "green",
+          "blue",
+          "indigo",
+          "violet",
+          "brown",
+          "black",
+          "slategray",
+         ]
+
+# Resized image width
+WIDTH = 600
+
+def detect(image):
+    width, height = image.size
+    ratio = float(WIDTH) / float(width)
+    new_h = height * ratio
+
+    image = image.resize((int(WIDTH), int(new_h)), Image.Resampling.LANCZOS)
+    
+    inputs = image_processor(images=image, return_tensors="pt")
+    outputs = model(**inputs)
+
+    # convert outputs to COCO API
+    target_sizes = torch.tensor([image.size[::-1]])
+    results = image_processor.post_process_object_detection(outputs,
+                                                            threshold=0.9,
+                                                            target_sizes=target_sizes)[0]
+
+    draw = ImageDraw.Draw(image)
+    
+    # label and the count
+    counts = {}
+
+    for score, label in zip(results["scores"], results["labels"]):
+        label_name = model.config.id2label[label.item()]
+        if label_name not in counts:
+            counts[label_name] = 0
+        counts[label_name] += 1
+
+    count_results = {k: v for k, v in (sorted(counts.items(), key=lambda item: item[1], reverse=True)[:10])}
+    label2color = {}
+    for idx, label in enumerate(count_results):
+        label2color[label] = colors[idx]
+
+    for label, box in zip(results["labels"], results["boxes"]):
+        label_name = model.config.id2label[label.item()]
+
+        if label_name in count_results:
+            box = [round(i, 4) for i in box.tolist()]
+            x1, y1, x2, y2 = tuple(box)
+            draw.rectangle((x1, y1, x2, y2), outline=label2color[label_name], width=2)
+            draw.text((x1, y1), label_name, fill="white")
+
+    df = pd.DataFrame({
+        'label': [label for label in count_results],
+        'counts': [counts[label] for label in count_results]
+    })
+    
+    return image, df, count_results
+
+demo = gr.Interface(
+    fn=detect,
+    examples=["examples/football.jpg", "examples/cats.jpg"],
+    inputs=[gr.inputs.Image(label="Input image", type="pil")],
+    outputs=[gr.Image(label="Output image"), gr.BarPlot(show_label=False, x="label", y="counts", x_title="Labels", y_title="Counts", vertical=False), gr.Textbox(show_label=False)],
+    title="YOLO Object Detection",
+    cache_examples=False
+)
+
+demo.launch()