Update app.py

app.py

@@ -3,17 +3,17 @@ from transformers import pipeline
 from PIL import Image, ImageDraw, ImageFont
 import tempfile
 
-# Load YOLOS object detection model
+# Load the YOLOS object detection model
 detector = pipeline("object-detection", model="hustvl/yolos-small")
 
-# Colors for labels
+# Define some colors to differentiate classes
 COLORS = ["red", "blue", "green", "orange", "purple", "yellow", "cyan", "magenta"]
 
-# Pick color based on label
+# Helper function to assign color per label
 def get_color_for_label(label):
     return COLORS[hash(label) % len(COLORS)]
 
-# Main detection + drawing function
+# Main function: detect, draw, and return outputs
 def detect_and_draw(image, threshold):
     results = detector(image)
     image = image.convert("RGB")
@@ -41,12 +41,7 @@ def detect_and_draw(image, threshold):
             width=3,
         )
 
-        draw.text(
-            (box["xmin"] + 5, box["ymin"] + 5),
-            label,
-            fill=color,
-            font=font
-        )
+        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))
@@ -55,9 +50,10 @@ def detect_and_draw(image, threshold):
     temp_file = tempfile.NamedTemporaryFile(suffix=".png", delete=False)
     image.save(temp_file.name)
 
-    return image, annotations, temp_file.name
+    # ✅ Return the (image, annotations) tuple and the path to the saved image
+    return (image, annotations), temp_file.name
 
-# Gradio UI
+# Gradio UI setup
 demo = gr.Interface(
     fn=detect_and_draw,
     inputs=[
@@ -68,8 +64,8 @@ demo = gr.Interface(
         gr.AnnotatedImage(label="Detected Image"),
         gr.File(label="Download Processed Image"),
     ],
-    title="YOLOS Object Detection (CPU Friendly)",
-    description="Upload an image to detect objects using the YOLOS-small model. Adjust the slider to set the confidence threshold.",
+    title="YOLOS Object Detection (CPU-Friendly)",
+    description="Upload an image to detect objects using the YOLOS-small model. Adjust the confidence threshold using the slider.",
 )
 
 demo.launch()