import os 
import numpy as np
from timeit import default_timer as timer
import cv2 as cv
import gradio as gr
from model_instance_function import get_pretrained_dog_emotion_classifier


# normalize function
def image_preprocessing(img):
  img = np.array(img)
  img = cv.resize(img,(224,224))
  img = img.reshape(1,224,224,3)
  return img / 255.0

# instance the model
model = get_pretrained_dog_emotion_classifier()

# gradio predict function
def predict(img):

  # class to map the indices to the classes
  class_2_index = {0: 'happy', 1: 'sad'}

  # measure execution time
  start_time = timer()

  # preprocess the image
  img = image_preprocessing(img)

  # make a prediction (prob of sad dog)
  pred_probability = model.predict(img)[0]

  # convert to an index
  pred_index = 1 if pred_probability > 0.5 else 0

  # label 
  pred_label = class_2_index[pred_index]

  end_time = timer()
  total_time = end_time - start_time

  return pred_probability, pred_label,round(total_time,5)


title = "Dog Emotions Vision Classifier"
description = "A vision classifier that distinguishes between sad and happy dogs."
article = "The model was trained in the [Dogs Emotions Dataset](https://huggingface.co/datasets/Q-b1t/Dogs_Emotions_Dataset) using the pretrained convolutional blocks of the VGG16 architecture and a custom classifier. For more information regarding the training, refer to this [colab notebook](https://colab.research.google.com/drive/1QqjLFsNV_8N1xr29BVwn4QVs_VH6lXmV?usp=sharing)."

example_list = [["examples/" + example] for example in os.listdir("examples")]

demo = gr.Interface(
    fn = predict,
    inputs = gr.Image(type = "pil"),
    outputs = [gr.Number(label = "Probability of a sad dog"),gr.Textbox(max_lines = 2,label = "Most likely class"),gr.Number(label = "Prediction time (s)")],
    examples = example_list,
    title = title,
    description = description,
    article = article
)

demo.launch()