|
import gradio as gr |
|
from fastai.learner import load_learner |
|
from fastai.vision.all import * |
|
import matplotlib.pyplot as plt |
|
|
|
|
|
def label_func(fname): |
|
categories = ["Happy", "Sad", "Angry"] |
|
category = fname.parts[-2] |
|
return category if category in categories else "unknown" |
|
|
|
learn = load_learner('emotion_model.pkl') |
|
|
|
|
|
def preprocess_image(image): |
|
|
|
image_pil = PILImage.create(image) |
|
|
|
image_resized = Resize(224)(image_pil) |
|
return image_resized |
|
|
|
|
|
def predict_emotion(image): |
|
|
|
processed_image = preprocess_image(image) |
|
|
|
pred_class, pred_idx, outputs = learn.predict(processed_image) |
|
|
|
class_names = learn.dls.vocab |
|
|
|
confidences = [outputs[idx].item() for idx in range(len(outputs))] |
|
|
|
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 |
|
|
|
|
|
iface = gr.Interface( |
|
fn=predict_emotion, |
|
inputs="image", |
|
outputs="plot", |
|
title="Emotion Prediction", |
|
description="Upload an image and get the confidence of predicted emotions" |
|
) |
|
|
|
|
|
iface.launch() |
|
|