# https://discuss.huggingface.co/t/how-to-fix-workload-evicted-storage-limit-exceeded-50g-error-in-huggingface-spaces/169258
import os
import torch

# CRITICAL: Redirect cache to temporary storage
os.environ['TORCH_HOME'] = '/tmp/torch_cache'
os.environ['HUB_DIR'] = '/tmp/torch_hub'
os.environ['TMPDIR'] = '/tmp'
torch.hub.set_dir('/tmp/torch_hub')

import glob
import gradio as gr
from ultralytics import YOLO

model_path = "best.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()