Update app.py

app.py

@@ -1,25 +1,69 @@
 import gradio as gr
 from transformers import pipeline
-from PIL import Image
+from PIL import Image, ImageDraw, ImageFont
 
-# Load YOLOS pipeline
+# Load YOLOS object detection model
 detector = pipeline("object-detection", model="hustvl/yolos-small")
 
-def detect_objects(img):
-    results = detector(img)
-    boxes = []
+# Confidence threshold (you can adjust this)
+CONFIDENCE_THRESHOLD = 0.5
+
+# Color palette for classes
+COLORS = [
+    "red", "blue", "green", "orange", "purple", "yellow", "cyan", "magenta"
+]
+
+def get_color_for_label(label):
+    return COLORS[hash(label) % len(COLORS)]
+
+def detect_and_draw(image):
+    results = detector(image)
+    
+    # Draw boxes on image
+    image = image.convert("RGB")
+    draw = ImageDraw.Draw(image)
+    
+    try:
+        font = ImageFont.truetype("arial.ttf", 16)
+    except:
+        font = ImageFont.load_default()
+
+    filtered_results = []
+    
     for obj in results:
-        label = f"{obj['label']} ({obj['score']:.2f})"
-        box = obj['box']
-        boxes.append((box["xmin"], box["ymin"], box["xmax"], box["ymax"], label))
-    return img, boxes
+        score = obj["score"]
+        if score < CONFIDENCE_THRESHOLD:
+            continue
+
+        label = f"{obj['label']} ({score:.2f})"
+        box = obj["box"]
+        color = get_color_for_label(obj["label"])
+
+        # Draw box
+        draw.rectangle(
+            [(box["xmin"], box["ymin"]), (box["xmax"], box["ymax"])],
+            outline=color,
+            width=3,
+        )
+        # Draw label
+        draw.text(
+            (box["xmin"] + 5, box["ymin"] + 5),
+            label,
+            fill=color,
+            font=font
+        )
+
+        # Also return data for Gradio AnnotatedImage
+        filtered_results.append((box["xmin"], box["ymin"], box["xmax"], box["ymax"], label))
+
+    return image, filtered_results
 
 demo = gr.Interface(
-    fn=detect_objects,
+    fn=detect_and_draw,
     inputs=gr.Image(type="pil"),
     outputs=gr.AnnotatedImage(),
     title="YOLOS Object Detection",
-    description="Upload an image and detect objects using the YOLOS Transformer model.",
+    description=f"Upload an image. Boxes show objects detected by the YOLOS Transformer model (confidence > {CONFIDENCE_THRESHOLD}).",
 )
 
 demo.launch()