import gradio as gr
from huggingface_hub import from_pretrained_fastai

# Carga del modelo desde HuggingFace Hub
learn = from_pretrained_fastai("ivferns/futurama")
labels = learn.dls.vocab

# Función de predicción
def classify_image(img):
    pred, idx, probs = learn.predict(img)
    # Mensaje de resultado
    message = f"La pelota que has subido es una {pred}"
    # Diccionario de probabilidades
    probs_dict = {labels[i]: float(probs[i]) for i in range(len(labels))}
    return message, probs_dict

# Interfaz Gradio
app = gr.Interface(
    fn=classify_image,
    inputs=gr.Image(type="pil", label="Sube una imagen de una pelota deportiva"),
    outputs=[
        gr.Textbox(label="Resultado"),
        gr.Label(num_top_classes=3, label="Probabilidades")
    ],
    title="Clasificador de Pelotas Deportivas",
    description="Sube una imagen y el modelo clasificará el tipo de pelota deportiva entre 15 categorías."
)

if __name__ == "__main__":
    app.launch(server_name="0.0.0.0", server_port=7860)