Update app.py

app.py

@@ -1,38 +1,65 @@
 import gradio as gr
-from transformers import YolosImageProcessor, YolosForObjectDetection
-from PIL import Image, ImageDraw
-import torch
+from transformers import pipeline
+from PIL import Image, ImageDraw, ImageFont
 
-# Load YOLOS processor and model
-processor = YolosImageProcessor.from_pretrained("hustvl/yolos-small")
-model = YolosForObjectDetection.from_pretrained("hustvl/yolos-small")
+# Load the YOLOS object detection model
+detector = pipeline("object-detection", model="hustvl/yolos-small")
 
-def detect_objects(image):
-    # Preprocess the image
-    inputs = processor(images=image, return_tensors="pt")
-    outputs = model(**inputs)
+# Define some colors to differentiate classes
+COLORS = ["red", "blue", "green", "orange", "purple", "yellow", "cyan", "magenta"]
 
-    # Post-process outputs
-    target_sizes = torch.tensor([image.size[::-1]])
-    results = processor.post_process_object_detection(outputs, target_sizes=target_sizes, threshold=0.5)[0]
+# Helper function to assign color per label
+def get_color_for_label(label):
+    return COLORS[hash(label) % len(COLORS)]
 
-    # Draw boxes on image
+# Main function: detect, draw, and return outputs
+def detect_and_draw(image, threshold):
+    results = detector(image)
+    image = image.convert("RGB")
     draw = ImageDraw.Draw(image)
-    for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
-        box = [round(i, 2) for i in box.tolist()]
-        label_text = f"{model.config.id2label[label.item()]}: {round(score.item(), 2)}"
-        draw.rectangle(box, outline="red", width=2)
-        draw.text((box[0], box[1]), label_text, fill="red")
-
-    return image
-
-# Gradio Interface
-iface = gr.Interface(
-    fn=detect_objects,
-    inputs=gr.Image(type="pil"),
-    outputs=gr.Image(type="pil"),
-    title="🟡 Object Detection with YOLOS",
-    description="Upload an image to detect objects using YOLOS (You Only Look One-level Series)."
+
+    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))
+
+    # Return the annotated image and annotations (no download option)
+    return image, annotations
+
+# 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.",
 )
 
-iface.launch()
+demo.launch()