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

model_id = "SEAR01/FER_model"
try:
    processor = AutoImageProcessor.from_pretrained(model_id)
    model = AutoModelForImageClassification.from_pretrained(model_id, trust_remote_code=True)  # 加這行允許自訂模型
except Exception as e:
    raise ValueError(f"Model load failed: {e}. Check repo files.")

emotion_labels = ['angry', 'disgust', 'fear', 'happy', 'neutral', 'sad', 'surprise']

def predict_emotion(image):
    if image is None:
        return "Upload an image."
    inputs = processor(images=image, return_tensors="pt")
    with torch.no_grad():
        outputs = model(**inputs)
        predicted = outputs.logits.argmax(-1).item()
    emotion = emotion_labels[predicted]
    confidence = torch.softmax(outputs.logits, dim=-1)[0][predicted].item()
    return f"Emotion: {emotion} (Confidence: {confidence:.2f})"

iface = gr.Interface(fn=predict_emotion, inputs=gr.Image(type="pil"), outputs="text", title="FER Demo")
if __name__ == "__main__":
    iface.launch()