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

# Cargar el modelo desde Hugging Face
learn = from_pretrained_fastai("abelllanas/emotions_dl")
labels = ["joy", "anger", "fear", "sadness"]  # Obtener las etiquetas de emociones

# Función para predecir emociones
def predict(img):
    img = PILImage.create(img)  # Convertir la imagen al formato adecuado
    pred, pred_idx, probs = learn.predict(img)  # Realizar la predicción
    return {labels[i]: float(probs[i]) for i in range(len(labels))} 


title = "Emotion predictor"
description = "A model based on paintings that tries to classify among joy, sadness, fear and anger."
examples = ['alegria_pintura.jpg', 'tristeza_pintura.jpg', 'miedo_pintura.jpg','ira_pintura.jpg']

gr.Interface(
    fn=predict,
    inputs=gr.Image(),
    outputs=gr.Label(num_top_classes=4),
    title=title,
    description=description,
    examples=examples,
).launch()