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):
    # Run inference
    results = model(image)

    # 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"),
    outputs=gr.Image(),
    title="Maritime Vessel Detection",
    description="Upload an image to detect vessels",
    examples=[["vessels.jpg"]],
    theme=gr.themes.Soft()
)

# Launch the app
demo.launch()