import gradio as gr
from fastai.vision.all import *
import skimage

learn = load_learner('export.pkl')

labels = learn.dls.vocab
def predict(img):
    img = PILImage.create(img)
    pred,pred_idx,probs = learn.predict(img)
    return {labels[i]: float(probs[i]) for i in range(len(labels))}

title = "Emotion Classifier"
description = "An emotion classifier trained on AffectNet with fastai. Identifies emotions of anger, contempt, disgust, fear, happy, neutral, surprise, and sad."

examples = ["happy.jpg", "neutral.jpg", "sad.jpg", "angry.jpg", "contempt.jpg", "fear.jpg", "surprise.jpg", "disgust.jpg"]

gr.Interface(
        fn=predict,
        inputs=gr.inputs.Image(shape=(224, 224)),
        outputs=gr.outputs.Label(num_top_classes=3),
        title=title,
        description=description,
        examples=examples,
        enable_queue=True,
).launch()