from pathlib import Path

import gradio as gr
from ultralytics import YOLO

# Load models
MODEL_PATH = "weights/mbari_315k_yolov8.pt"
MODEL = YOLO(MODEL_PATH)

PREDICT_KWARGS = {
    "conf": 0.15,
}

# Get example images
EXAMPLES_DIR = Path("examples")
EXAMPLES = list(EXAMPLES_DIR.glob("*.png")) if EXAMPLES_DIR.exists() else []


def detect_objects(image):
    results = MODEL.predict(image, **PREDICT_KWARGS)
    return results[0].plot()


# Gradio interface with gr.Interface instead of gr.Blocks
demo = gr.Interface(
    fn=detect_objects,
    inputs=gr.Image(type="numpy"),
    outputs=gr.Image(type="numpy"),
    title="MBARI 315k",
    examples=EXAMPLES if EXAMPLES else None,
    cache_examples=True,
)

demo.queue().launch()