SakibRumu
Update app.py
1675d25 verified
raw
history blame
2 kB
import gradio as gr
from fastai.vision.all import load_learner, PILImage
from PIL import Image
import torch
# Define the class mapping for the RAF-DB dataset
class_mapping = {
"1": "Surprise",
"2": "Fear",
"3": "Disgust",
"4": "Happiness",
"5": "Sadness",
"6": "Anger",
"7": "Contempt"
}
# Define the function to map folder names to labels
def get_raf_label(file_path):
# Use the class mapping to get the label
return class_mapping[str(file_path.parent.name)]
# Load the model
model_path = "efficientnet_emotion_model.pkl" # Replace with your actual model path
model = load_learner(model_path)
# Define the emotion classes
emotion_classes = list(class_mapping.values()) # Get emotion classes from the class mapping
# Function for Emotion Prediction
def predict_emotion(image):
img = PILImage.create(image) # Convert the uploaded image into a PIL image
pred_class, pred_idx, outputs = model.predict(img)
predicted_emotion = emotion_classes[pred_idx]
confidence = outputs[pred_idx] * 100 # Convert to percentage
return predicted_emotion, f"{confidence:.2f}%"
# Gradio interface with xkcd theme
with gr.Blocks(theme="gstaff/xkcd") as demo:
gr.Markdown("# Emotion Recognition Classifier")
gr.Markdown("""
This app uses a deep learning model to recognize emotions in facial images.
The model has been trained on a dataset to classify images into different emotion categories:
* Anger
* Fear
* Happiness
* Sadness
* Surprise
* Neutral
""")
# Upload image widget
image_input = gr.Image(type="pil", label="Upload an image of a face")
# Outputs
label_output = gr.Textbox(label="Predicted Emotion")
confidence_output = gr.Textbox(label="Confidence Percentage")
# Button to predict the emotion
image_input.upload(predict_emotion, image_input, [label_output, confidence_output])
# Launch the app
demo.launch(share=True)