Create app.py

app.py

@@ -0,0 +1,71 @@
+import gradio as gr
+from transformers import pipeline
+from PIL import Image, ImageDraw, ImageFont
+import tempfile
+
+# 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):
+    results = detector(image)
+    image = image.convert("RGB")
+    draw = ImageDraw.Draw(image)
+
+    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.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))
+
+    # Save image for download
+    temp_file = tempfile.NamedTemporaryFile(suffix=".png", delete=False)
+    image.save(temp_file.name)
+
+    # ✅ Return the (image, annotations) tuple and the path to the saved image
+    return (image, annotations), temp_file.name
+
+# 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"),
+        gr.File(label="Download Processed Image"),
+    ],
+    title="YOLOS Object Detection",
+    description="Upload an image to detect objects using the YOLOS-small model. Adjust the confidence threshold using the slider.",
+)
+
+demo.launch()