import gradio as gr

# For loading files
from joblib import dump, load

# Model hub
import tensorflow_hub as hub

# Neural networks
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers
from keras.applications.vgg16 import preprocess_input
from huggingface_hub import from_pretrained_keras

# Image processing
import PIL

#------------------------------------------

# Loading model
model = from_pretrained_keras('ana-bernal/keras_dog_breed_eff')

# Reading file with class names, uncomment to import class names
breed_names_norm = []
with open('labels.txt', 'r') as file:
    for line in file:
        # remove linebreak from a current name
        name = line[:-1]
        breed_names_norm.append(name)

# Definition of main function
def classify_image(inp):
    """
    Returns a dictionnary: predicted_breeds, where the
    keys are [1,2,3] for the first, second and third more probable
    breed for the dog image. Each value is a dictionnary with keys
    ['idx', 'name', 'confidence'] and their corresponding values.
    
    Parameters:
        img: returned by the function load_img_path
    """
    img_array = keras.preprocessing.image.img_to_array(inp)
    img_array = tf.expand_dims(img_array, 0)  # Creates a batch axis

    predictions = model.predict(img_array, verbose=0).flatten()
    confidences = {breed_names_norm[i]: float(predictions[i]) for i in range(120)}

    return confidences

# --------------------------------------------------

examples = [
    ['example_images/01_test.jpg'],
    ['example_images/02_test.jpg'],
    ['example_images/03_test.jpg'],
    ['example_images/04_test.jpg'],
    ['example_images/05_test.jpg'],
    ['example_images/06_test.jpg'],
]

demo = gr.Interface(fn=classify_image, 
                    inputs=gr.Image(shape=(180, 180)), 
                    outputs=gr.Label(num_top_classes=3),
                    examples=examples)


if __name__ == "__main__":
    demo.launch()