Create app.py

app.py

@@ -0,0 +1,72 @@
+import gradio as gr
+from transformers import pipeline
+from PIL import Image, ImageDraw, ImageFont
+
+# Load the YOLOS object detection model
+detector = pipeline("object-detection", model="hustvl/yolos-small")
+
+# Define some colors to differentiate classes
+COLORS = ["red", "blue", "green", "orange", "purple", "yellow", "cyan", "magenta"]
+
+# Helper function to assign color per label
+def get_color_for_label(label):
+    return COLORS[hash(label) % len(COLORS)]
+
+# Main function: detect, draw, and return outputs
+def detect_and_draw(image, threshold):
+    try:
+        # Perform object detection
+        results = detector(image)
+        image = image.convert("RGB")
+        draw = ImageDraw.Draw(image)
+
+        # Try to load a font for annotations, else use default
+        try:
+            font = ImageFont.truetype("arial.ttf", 16)
+        except:
+            font = ImageFont.load_default()
+
+        annotations = []
+
+        for obj in results:
+            score = obj["score"]
+            if score < threshold:
+                continue
+
+            label = f"{obj['label']} ({score:.2f})"
+            box = obj["box"]
+            color = get_color_for_label(obj["label"])
+
+            # Draw the bounding box and label
+            draw.rectangle(
+                [(box["xmin"], box["ymin"]), (box["xmax"], box["ymax"])],
+                outline=color,
+                width=3,
+            )
+
+            draw.text((box["xmin"] + 5, box["ymin"] + 5), label, fill=color, font=font)
+
+            box_coords = (box["xmin"], box["ymin"], box["xmax"], box["ymax"])
+            annotations.append((box_coords, label))
+
+        return image, annotations
+
+    except Exception as e:
+        return f"Error during detection: {e}", None
+
+# Gradio UI setup
+demo = gr.Interface(
+    fn=detect_and_draw,
+    inputs=[
+        gr.Image(type="pil", label="Upload Image"),
+        gr.Slider(minimum=0.1, maximum=1.0, value=0.5, step=0.05, label="Confidence Threshold"),
+    ],
+    outputs=[
+        gr.AnnotatedImage(label="Detected Image"),
+    ],
+    title="YOLOS Object Detection",
+    description="Upload an image to detect objects using the YOLOS-small model. Adjust the confidence threshold using the slider.",
+    live=True
+)
+
+demo.launch()