| 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() | |