# Facial expression classifier
import os
from fastai.vision.all import *
import gradio as gr

# Emotion
learn_emotion = load_learner('emotions_vgg19.pkl')
learn_emotion_labels = learn_emotion.dls.vocab

# Sentiment
learn_sentiment = load_learner('sentiment_vgg19.pkl')
learn_sentiment_labels = learn_sentiment.dls.vocab

# Predict
def predict(img):
    img = PILImage.create(img)
    
    pred_emotion, pred_emotion_idx, probs_emotion = learn_emotion.predict(img)
    
    pred_sentiment, pred_sentiment_idx, probs_sentiment = learn_sentiment.predict(img)
    
    #emotions = {f'emotion_{learn_emotion_labels[i]}': float(probs_emotion[i]) for i in range(len(learn_emotion_labels))}
    #sentiments = {f'sentiment_{learn_sentiment_labels[i]}': float(probs_sentiment[i]) for i in range(len(learn_sentiment_labels))}
    
    emotions = {learn_emotion_labels[i]: float(probs_emotion[i]) for i in range(len(learn_emotion_labels))}
    sentiments = {learn_sentiment_labels[i]: float(probs_sentiment[i]) for i in range(len(learn_sentiment_labels))}
        
    return [emotions, sentiments] #{**emotions, **sentiments}

# Gradio
title = "Facial Emotion and Sentiment Detector"

description = gr.Markdown(
                """Ever wondered what a person might be feeling looking at their picture? 
                 Well, now you can! Try this fun app. Just upload a facial image in JPG or
                 PNG format. Voila! you can now see what they might have felt when the picture
                 was taken.
                 
                 **Tip**: Be sure to only include face to get best results. Check some sample images
                 below for inspiration!""").value

article = gr.Markdown(
             """**DISCLAIMER:** This model does not reveal the actual emotional state of a person. Use and 
             
             
             Positive (Happy, Surprise)
             
             Negative (Angry, Disgust, Fear, Sad)
             
             Neutral (Neutral)
             
             **MODEL:** VGG19""").value

enable_queue=True

examples = ['happy1.jpg', 'happy2.jpg', 'angry1.png', 'angry2.jpg', 'neutral1.jpg', 'neutral2.jpg']

gr.Interface(fn = predict, 
             inputs = gr.Image( image_mode='L'), 
             outputs = [gr.Label(label='Emotion'), gr.Label(label='Sentiment')], #gr.Label(),
             title = title,
             examples = examples,
             description = description,
             article=article,
             allow_flagging='never').launch()