import gradio as gr
from transformers import AutoModel, AutoFeatureExtractor, AutoConfig
import torch
from PIL import Image

# Carga el modelo y el extractor de características
model_name = "CarPeAs/reconocimiento-facial"

# Intentar cargar la configuración del modelo primero
config = AutoConfig.from_pretrained(model_name)

# Imprimir la configuración del modelo para verificar su contenido
print(config)

# Usar la clase correcta basada en la configuración del modelo
if config.model_type in ['vit', 'beit', 'deit', 'resnet', 'convnext', 'swin']:
    model = AutoModelForImageClassification.from_pretrained(model_name)
else:
    model = AutoModel.from_pretrained(model_name)

feature_extractor = AutoFeatureExtractor.from_pretrained(model_name)

# Función de inferencia
def reconocer_emocion(image):
    # Convierte la imagen a tensor
    inputs = feature_extractor(images=image, return_tensors="pt")
    with torch.no_grad():
        outputs = model(**inputs)
    # Asumiendo que las emociones están en las etiquetas de salida del modelo
    logits = outputs.logits
    emotion = torch.argmax(logits, dim=-1).item()
    emociones = ["Feliz", "Triste", "Enojado", "Sorprendido", "Neutral"]
    return emociones[emotion]

# Interfaz de Gradio
iface = gr.Interface(
    fn=reconocer_emocion,
    inputs=gr.inputs.Image(type="pil"),
    outputs=gr.outputs.Textbox(label="Emoción Reconocida"),
    title="Reconocimiento Facial de Emociones",
    description="Carga una imagen y el modelo reconocerá la emoción expresada en el rostro."
)

# Ejecuta la aplicación
if __name__ == "__main__":
    iface.launch()