VIto100's picture
Create app.py
1e9d2c9 verified
import gradio as gr
from transformers import BeitImageProcessor, BeitForImageClassification
from PIL import Image
# Load processor and model globally to avoid reloading on each function call
processor = BeitImageProcessor.from_pretrained("Tanneru/Facial-Emotion-Detection-FER-RAFDB-AffectNet-BEIT-Large")
model = BeitForImageClassification.from_pretrained("Tanneru/Facial-Emotion-Detection-FER-RAFDB-AffectNet-BEIT-Large")
emotion_map = {
"LABEL_0": "Angry ๐Ÿ˜ ",
"LABEL_1": "Disgust ๐Ÿคข",
"LABEL_2": "Fear ๐Ÿ˜จ",
"LABEL_3": "Happy ๐Ÿ˜Š",
"LABEL_4": "Sad ๐Ÿ˜”",
"LABEL_5": "Surprise ๐Ÿ˜ฎ",
"LABEL_6": "Neutral ๐Ÿ˜"
}
def predict_emotion(image_numpy):
if image_numpy is None:
return "Please upload an image."
# Convert numpy image (from Gradio) to PIL Image
image_pil = Image.fromarray(image_numpy).convert("RGB")
inputs = processor(images=image_pil, return_tensors="pt")
outputs = model(**inputs)
predicted_class = outputs.logits.argmax(-1).item()
generic_label = model.config.id2label[predicted_class]
descriptive_emotion = emotion_map.get(generic_label, generic_label)
return f"Predicted emotion: {descriptive_emotion}"
# Create and launch the Gradio interface
interface = gr.Interface(
fn=predict_emotion,
inputs=gr.Image(type="numpy", label="Input Image"),
outputs="text",
title="Facial Emotion Detection"
)
interface.launch()