File size: 1,406 Bytes
1e9d2c9
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
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()