import gradio as gr
import os
import numpy as np
import tensorflow as tf
import keras
import keras_cv
from keras.models import load_model
import cv2


def image_predict(img_):
    model = load_model('efficientnet_b0.keras')
    # img = cv2.cvtColor(img_, cv2.COLOR_BGR2RGB)
    img = cv2.resize(img_, dsize = [224, 224])
    img = img / 255.0
    img = np.expand_dims(img, axis = 0)

    pred = model.predict(img, verbose = 1)
    pred = np.argmax(pred, axis = 1)

    classes = ['angry', 'happy', 'neutral', 'sad', 'suprised', 'tired']
    if pred == 0:
        answer = f"Facial Expression detected is: {classes[0].capitalize()}"
    elif pred == 1:
        answer = f"Facial Expression detected is: {classes[1].capitalize()}"
    elif pred == 2:
        answer = f"Facial Expression detected is: {classes[2].capitalize()}"
    elif pred == 3:
        answer = f"Facial Expression detected is: Depressed"
    elif pred == 4:
        answer = f"Facial Expression detected is: {classes[4].capitalize()}"
    elif pred == 5:
        answer = f"Facial Expression detected is: {classes[5].capitalize()}"

    return answer


with gr.Blocks() as demo:
    image_ = gr.Image(label = 'Input Image to be predicted')
    output = gr.Textbox(label = 'Prediction')
    btn = gr.Button('Predict')
    btn.click(fn = image_predict, inputs = [image_], outputs = output)

demo.launch(share = False)