import glob
import gradio as gr
from ultralytics import YOLO

model_path = "fathomnet23-comp-baseline.pt"
model = YOLO(model_path)


def run(image_path):
    results = model.predict(image_path)
    return results[0].plot()[:, :, ::-1]  # reverse channels for gradio


title = "FathomNet2023 Competition Baseline"
description = (
    "Gradio demo for the FathomNet2023 Baseline Model: Developed by researchers"
    " at the Monterey Bay Aquarium Research Institute (MBARI) to serve as a"
    " baseline YOLOv8m model for the FathomNet2023 Kaggle Competition, in"
    " conjunction with the Fine Grained Visual Categorization workshop at CVPR"
    " 2023. The training dataset comprises both the FathomNet2023 competition"
    " split and internal MBARI data, including 290 fine-grained taxonomic"
    " categories of benthic animals."
)

examples = glob.glob("images/*.png")

interface = gr.Interface(
    run,
    inputs=[gr.components.Image(type="filepath")],
    outputs=gr.components.Image(type="numpy"),
    title=title,
    description=description,
    examples=examples,
)

interface.queue().launch()