import gradio as gr
from huggingface_hub import hf_hub_download
from ultralytics import YOLO
import cv2

# Download and load the model
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, 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)

# Create Gradio interface
demo = gr.Interface(
    fn=detect_vessels,
    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. Adjust confidence and IoU thresholds to filter detections.",
    examples=[["vessels.jpg", 0.25, 0.7]],
    theme=gr.themes.Soft()
)

# Launch the app
demo.launch()