Update app.py

app.py

@@ -1,23 +1,21 @@
 import gradio as gr
 from transformers import pipeline
 from PIL import Image, ImageDraw, ImageFont
+import tempfile
 
 # Load YOLOS object detection model
 detector = pipeline("object-detection", model="hustvl/yolos-small")
 
-# Confidence threshold
-CONFIDENCE_THRESHOLD = 0.5
-
-# Color palette
+# Colors for labels
 COLORS = ["red", "blue", "green", "orange", "purple", "yellow", "cyan", "magenta"]
 
+# Pick color based on label
 def get_color_for_label(label):
     return COLORS[hash(label) % len(COLORS)]
 
-def detect_and_draw(image):
+# Main detection + drawing function
+def detect_and_draw(image, threshold):
     results = detector(image)
-    
-    # Convert to RGB for drawing
     image = image.convert("RGB")
     draw = ImageDraw.Draw(image)
 
@@ -30,21 +28,19 @@ def detect_and_draw(image):
 
     for obj in results:
         score = obj["score"]
-        if score < CONFIDENCE_THRESHOLD:
+        if score < 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,
@@ -52,19 +48,28 @@ def detect_and_draw(image):
             font=font
         )
 
-        # AnnotatedImage expects (box_tuple, label)
         box_coords = (box["xmin"], box["ymin"], box["xmax"], box["ymax"])
         annotations.append((box_coords, label))
 
-    return image, annotations
+    # Save image for download
+    temp_file = tempfile.NamedTemporaryFile(suffix=".png", delete=False)
+    image.save(temp_file.name)
+
+    return image, annotations, temp_file.name
 
-# Gradio interface
+# Gradio UI
 demo = gr.Interface(
     fn=detect_and_draw,
-    inputs=gr.Image(type="pil"),
-    outputs=gr.AnnotatedImage(),
-    title="YOLOS Object Detection",
-    description=f"Upload an image to detect objects using the YOLOS model. Only objects with confidence > {CONFIDENCE_THRESHOLD} are shown.",
+    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 (CPU Friendly)",
+    description="Upload an image to detect objects using the YOLOS-small model. Adjust the slider to set the confidence threshold.",
 )
 
 demo.launch()