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