Update app.py

app.py

@@ -5,13 +5,11 @@ from PIL import Image, ImageDraw, ImageFont
 # Load YOLOS object detection model
 detector = pipeline("object-detection", model="hustvl/yolos-small")
 
-# Confidence threshold (you can adjust this)
+# Confidence threshold
 CONFIDENCE_THRESHOLD = 0.5
 
-# Color palette for classes
-COLORS = [
-    "red", "blue", "green", "orange", "purple", "yellow", "cyan", "magenta"
-]
+# Color palette
+COLORS = ["red", "blue", "green", "orange", "purple", "yellow", "cyan", "magenta"]
 
 def get_color_for_label(label):
     return COLORS[hash(label) % len(COLORS)]
@@ -19,17 +17,17 @@ def get_color_for_label(label):
 def detect_and_draw(image):
     results = detector(image)
     
-    # Draw boxes on image
+    # Convert to RGB for drawing
     image = image.convert("RGB")
     draw = ImageDraw.Draw(image)
-    
+
     try:
         font = ImageFont.truetype("arial.ttf", 16)
     except:
         font = ImageFont.load_default()
 
-    filtered_results = []
-    
+    annotations = []
+
     for obj in results:
         score = obj["score"]
         if score < CONFIDENCE_THRESHOLD:
@@ -45,6 +43,7 @@ def detect_and_draw(image):
             outline=color,
             width=3,
         )
+
         # Draw label
         draw.text(
             (box["xmin"] + 5, box["ymin"] + 5),
@@ -53,17 +52,19 @@ def detect_and_draw(image):
             font=font
         )
 
-        # Also return data for Gradio AnnotatedImage
-        filtered_results.append((box["xmin"], box["ymin"], box["xmax"], box["ymax"], label))
+        # AnnotatedImage expects (box_tuple, label)
+        box_coords = (box["xmin"], box["ymin"], box["xmax"], box["ymax"])
+        annotations.append((box_coords, label))
 
-    return image, filtered_results
+    return image, annotations
 
+# Gradio interface
 demo = gr.Interface(
     fn=detect_and_draw,
     inputs=gr.Image(type="pil"),
     outputs=gr.AnnotatedImage(),
     title="YOLOS Object Detection",
-    description=f"Upload an image. Boxes show objects detected by the YOLOS Transformer model (confidence > {CONFIDENCE_THRESHOLD}).",
+    description=f"Upload an image to detect objects using the YOLOS model. Only objects with confidence > {CONFIDENCE_THRESHOLD} are shown.",
 )
 
 demo.launch()