Update app.py

app.py

@@ -8,10 +8,10 @@ repo_id = "truthdotphd/vessel-detection"
 model_path = hf_hub_download(repo_id=repo_id, filename="model.pt")
 model = YOLO(model_path)
 
-def detect_vessels(image):
-    # Run inference
-    results = model(image)
-
+def detect_vessels(image, conf_threshold, iou_threshold):
+    # Run inference with custom thresholds
+    results = model(image, conf=conf_threshold, iou=iou_threshold)
+    
     # Plot results
     annotated_image = results[0].plot()
     return cv2.cvtColor(annotated_image, cv2.COLOR_BGR2RGB)
@@ -19,13 +19,17 @@ def detect_vessels(image):
 # Create Gradio interface
 demo = gr.Interface(
     fn=detect_vessels,
-    inputs=gr.Image(type="numpy"),
+    inputs=[
+        gr.Image(type="numpy"),
+        gr.Slider(minimum=0.1, maximum=1.0, value=0.25, step=0.05, label="Confidence Threshold"),
+        gr.Slider(minimum=0.1, maximum=1.0, value=0.7, step=0.05, label="IoU Threshold")
+    ],
     outputs=gr.Image(),
     title="Maritime Vessel Detection",
-    description="Upload an image to detect vessels",
-    examples=[["vessels.jpg"]],
+    description="Upload an image to detect vessels. Adjust confidence and IoU thresholds to filter detections.",
+    examples=[["vessels.jpg", 0.25, 0.7]],
     theme=gr.themes.Soft()
 )
 
 # Launch the app
-demo.launch()
+demo.launch()