import gradio as gr
from fastai.vision.all import *

# Carga tu modelo
learn = load_learner("enemies_BOTW.pkl")  # Reemplaza con la ruta a tu modelo

# Define la función de predicción
def predict(img):
    img = PILImage.create(img)
    pred, pred_idx, probs = learn.predict(img)
    return {learn.dls.vocab[i]: float(probs[i]) for i in range(len(learn.dls.vocab))}

# Crea la interfaz de Gradio
gr.Interface(
    fn=predict,
    inputs=gr.Image(),
    outputs=gr.Label(num_top_classes=2),
    title="Clasificador de Enemigos de BOTW",
    description = "Sube una imagen de un enemigo del Zelda BOTW para clasificarlo. Solo puede clasificar entre: bokoblins, lizalfos, chuchus, keese y centaleones",
    examples=[
        ["2c1fe44c-105b-4c16-91d0-21dc58b53d6d.jpg"],  
        ["Captura de pantalla de 2024-11-14 19-01-45.png"],
        ["Captura de pantalla de 2024-12-21 18-02-26.png"],
        ["Captura de pantalla de 2024-12-21 21-39-33.png"],
        ["IMG_2349.jpg"] 
    ]

).launch()