# -*- coding: utf-8 -*-
"""gradio_deploy.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/13X2E9v7GxryXyT39R5CzxrNwxfA6KMFJ
"""
import os
import gradio as gr
from PIL import Image
from timeit import default_timer as timer
from tensorflow import keras
import numpy as np

MODEL = keras.models.load_model("convnet_from_scratch_with_augmentation.keras")

def predict(img):

    # Start the timer
    start_time = timer()

    # Reading the image and size transformation
    features = Image.open(img)
    features = features.resize((180, 180))
    features = np.array(features).reshape(1, 180,180,3)

    # Create a prediction label and prediction probability dictionary for each prediction class
    # This is the required format for Gradio's output parameter
    pred_labels_and_probs = {'dog':float(MODEL.predict(features))} if MODEL.predict(features)> 0.5 else {'cat':1-float(MODEL.predict(features))}

    # Calculate the prediction time
    pred_time = round(timer() - start_time, 5)

    # Return the prediction dictionary and prediction time
    return pred_labels_and_probs, pred_time

# Create title, description and article strings
example_list = [["examples/" + example] for example in os.listdir("examples")]
title = "Classification Demo"
description = "Cat/Dog classification Tensorflow model with Augmented small dataset"

# Create the Gradio demo
demo = gr.Interface(fn=predict, # mapping function from input to output
                    inputs=gr.Image(type='filepath'), # what are the inputs?
                    outputs=[gr.Label(label="Predictions"), # what are the outputs?
                             gr.Number(label="Prediction time (s)")], # our fn has two outputs, therefore we have two outputs
                    examples=example_list,
                    title=title,
                    description=description,)

# Launch the demo!
demo.launch()