# AUTOGENERATED! DO NOT EDIT! File to edit: ../app.ipynb.

# %% auto 0
__all__ = ['device', 'model', 'CLASS_LABELS', 'image', 'label', 'examples', 'intf', 'classify_emotions']

# %% ../app.ipynb 2
import gradio as gr 
import torch
from torch.nn.functional import softmax
import numpy as np

# %% ../app.ipynb 3
device = "cuda" if torch.cuda.is_available() else "cpu"

model = torch.load('model.pth').to(device)
model.eval()

# %% ../app.ipynb 4
CLASS_LABELS  = ['Anger', 'Disgust', 'Fear', 'Happy', 'Neutral', 'Sadness', "Surprise"]

# %% ../app.ipynb 5
def classify_emotions(im):
    im = np.array(im) / 255
    if len(im.shape) == 2:
        im = im[..., np.newaxis]
    if im.shape[-1] == 1:
        im = np.concatenate((im, im, im), 2)
    im = torch.tensor(im.transpose(2, 0, 1), dtype=torch.float32)
    prediction = model.forward(im[np.newaxis, ...].to(device))
    return dict(zip(CLASS_LABELS, *softmax(prediction, dim=1).tolist()))
    

# %% ../app.ipynb 6
image = gr.inputs.Image((48, 48))
label = gr.outputs.Label()
examples = ['happy.png', 'fear.png', 'anger.png']

intf = gr.Interface(fn=classify_emotions, 
                    inputs=image, 
                    outputs=label, 
                    title='Emotion classification',
                    examples=examples)
intf.launch(inline=False)