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

# ✅ Ganti label sesuai dataset kamu
id2label = {
    0: "Angry",
    1: "Disgust",
    2: "Fear",
    3: "Happy",
    4: "Sad",
    5: "Surprise",
    6: "Neutral"
}

# ✅ Ganti checkpoint/model sesuai modelmu
checkpoint = "google/vit-base-patch16-224-in21k"
processor = AutoImageProcessor.from_pretrained(checkpoint)
model = AutoModelForImageClassification.from_pretrained(
    checkpoint,
    num_labels=7,
    id2label=id2label,
    label2id={v: k for k, v in id2label.items()}
)
model.eval()

def predict(img: Image.Image):
    if img.mode != "RGB":
        img = img.convert("RGB")
    inputs = processor(images=img, return_tensors="pt").to(model.device)
    with torch.no_grad():
        outputs = model(**inputs)
        logits = outputs.logits
        probs = torch.nn.functional.softmax(logits, dim=1)[0]
    return {id2label[i]: float(probs[i]) for i in range(len(probs))}

interface = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil"),
    outputs=gr.Label(num_top_classes=3),
    title="Facial Expression Classifier",
    description="Upload a face image and the model will predict the facial expression."
)

if __name__ == "__main__":
    interface.launch()