import os

# Keep framework caches on ephemeral scratch space rather than the Space's
# persistent disk, to avoid exceeding the storage quota:
# https://discuss.huggingface.co/t/how-to-fix-workload-evicted-storage-limit-exceeded-50g-error-in-huggingface-spaces/169258
os.environ.setdefault("TORCH_HOME", "/tmp/torch_cache")
os.environ.setdefault("HUB_DIR", "/tmp/torch_hub")
os.environ.setdefault("HF_HOME", "/tmp/hf_home")
os.environ.setdefault("TMPDIR", "/tmp")

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

MODEL_PATH = hf_hub_download(
    repo_id="FathomNet/megalodon-2023-yolov8",
    filename="mbari-megalodon-yolov8x.pt",
)
model = YOLO(MODEL_PATH)


PREDICT_KWARGS = {
    "classes": 0,
    "conf": 0.25,
}


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


title = "Megalodon Detector"
description = (
    ""
)

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()