import json
import tensorflow as tf
from huggingface_hub import hf_hub_download
import gradio as gr

tf_model = hf_hub_download(repo_id='mikachou/dog-breed-classifier', filename='tf_model.h5')
config_json = hf_hub_download(repo_id='mikachou/dog-breed-classifier', filename='config.json')

model = tf.keras.models.load_model(tf_model)
print(model.summary())

with open(config_json) as f:
    config = json.load(f)

dogs_breeds = list(config['id2label'].values())

def predict(filepath):
    img = tf.io.read_file(filepath)
    tensor = tf.io.decode_image(img, channels=3, dtype=tf.dtypes.float32)
    tensor = tf.image.resize(tensor, [299, 299])
    input_tensor = tf.expand_dims(tensor, axis=0)

    output = model.predict(input_tensor)

    confidences = { dogs_breeds[i]: float(output[0][i]) for i in range(120) }

    return confidences

demo = gr.Interface(
    fn=predict,
    inputs=gr.Image(label='photo', type='filepath'),
    outputs=gr.Label(label="Predicted breed", num_top_classes=3),
    examples=[
        'imgs/beethoven.jpg',
        'imgs/belle.png',
        'imgs/belmondo.jpg',
        'imgs/dorothy.jpg',
        'imgs/lassie.jpg',
        'imgs/rintintin.jpg'
    ],
    title="Dog breed identification",
    description="The model was trained with [Stanford Dogs Dataset](http://vision.stanford.edu/aditya86/ImageNetDogs/) using tensorflow/keras on a fine-tuned pre-trained InceptionResNetV2 model",
    article="You could also drag/drop other examples from [this page](https://www.rdasia.com/pets/can-you-guess-dog-breed-based-its-puppy-picture)")

demo.launch()