Orphen2024's picture
Update app.py
18dde07 verified
raw
history blame
1.62 kB
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()