Update app.py

app.py

@@ -1,33 +1,40 @@
 import gradio as gr
 from ultralytics import YOLO
-import cv2
 import numpy as np
 
-# Load your YOLO11n model
+# Load YOLO11n model
 model = YOLO("yolo11n.pt")
 
-def detect_obstacles(image):
-    """
-    Run YOLO11n model on the uploaded image and return detections.
-    """
-    # Convert PIL image to numpy
+def detect_obstacles(image, conf_threshold):
     img = np.array(image)
 
-    # Run inference
-    results = model.predict(img)
+    # Run inference with confidence threshold
+    results = model.predict(img, conf=conf_threshold)
 
-    # Draw detections on image
-    annotated_img = results[0].plot()  # Ultralytics provides a plot() method
+    # Annotated image
+    annotated_img = results[0].plot()
 
-    return annotated_img
+    # Extract labels for debugging
+    labels = []
+    for box in results[0].boxes:
+        cls_id = int(box.cls[0])
+        conf = float(box.conf[0])
+        labels.append(f"{model.names[cls_id]} ({conf:.2f})")
+
+    return annotated_img, "\n".join(labels)
 
-# Gradio interface
 demo = gr.Interface(
     fn=detect_obstacles,
-    inputs=gr.Image(type="pil", label="Upload an Image"),
-    outputs=gr.Image(type="numpy", label="Detected Obstacles"),
+    inputs=[
+        gr.Image(type="pil", label="Upload Image"),
+        gr.Slider(0, 1, value=0.5, label="Confidence Threshold")
+    ],
+    outputs=[
+        gr.Image(type="numpy", label="Detected Obstacles"),
+        gr.Textbox(label="Detections")
+    ],
     title="YOLO11n Obstacle Detection",
-    description="Upload an image and the YOLO11n model will detect obstacles."
+    description="Upload an image. Adjust confidence threshold to filter detections."
 )
 
 if __name__ == "__main__":