emotion / app.py
RoAr777's picture
Update app.py
1ad01ef verified
raw
history blame
1.58 kB
import gradio as gr
from fastai.learner import load_learner
from fastai.vision.all import *
import matplotlib.pyplot as plt
# Load the trained model
def label_func(fname):
categories = ["Happy", "Sad", "Angry"]
category = fname.parts[-2] # Extract the category name from the path
return category if category in categories else "unknown"
learn = load_learner('emotion_model.pkl')
# Define function to preprocess image
def preprocess_image(image):
# Convert image to PILImage
image_pil = PILImage.create(image)
# Apply same transformations as done during training
image_resized = Resize(224)(image_pil)
return image_resized
# Define function to predict emotion and generate confidence bar chart
def predict_emotion(image):
# Preprocess the image
processed_image = preprocess_image(image)
# Predict the emotion
pred_class, pred_idx, outputs = learn.predict(processed_image)
# Get class names
class_names = learn.dls.vocab
# Generate confidence values
confidences = [outputs[idx].item() for idx in range(len(outputs))]
# Generate bar chart
plt.bar(class_names, confidences)
plt.xlabel('Emotion')
plt.ylabel('Confidence')
plt.title('Confidence of Predicted Emotions')
plt.xticks(rotation=45)
plt.tight_layout()
return plt
# Define Gradio interface
iface = gr.Interface(
fn=predict_emotion,
inputs="image",
outputs="plot",
title="Emotion Prediction",
description="Upload an image and get the confidence of predicted emotions"
)
# Launch the interface
iface.launch()