import gradio as gr
from transformers import AutoImageProcessor, AutoModelForImageClassification
import torch
from PIL import Image

# Load model and processor
processor = AutoImageProcessor.from_pretrained("dima806/facial_emotions_image_detection")
model = AutoModelForImageClassification.from_pretrained("dima806/facial_emotions_image_detection")

# Function to predict emotion
def detect_emotion(image):
    inputs = processor(images=image, return_tensors="pt")
    with torch.no_grad():
        outputs = model(**inputs)
    logits = outputs.logits
    predicted_class = logits.argmax(-1).item()
    return model.config.id2label[predicted_class]

# Gradio interface
iface = gr.Interface(
    fn=detect_emotion,
    inputs=gr.Image(type="pil"),
    outputs="text",
    title="Facial Emotion Detection",
    description="Upload a face image to detect emotion using dima806 model"
)

iface.launch(share=True)