# Workaround to install the lib without "setup.py"
import sys
from git import Repo
Repo.clone_from("https://github.com/dimitreOliveira/hub.git", "./hub")
sys.path.append("/hub")

import gradio as gr
from hub.tensorflow_hub.hf_utils import pull_from_hub

import requests
# Download human-readable labels for ImageNet.
response = requests.get("https://storage.googleapis.com/download.tensorflow.org/data/ImageNetLabels.txt")
labels = [x for x in response.text.split("\n") if x != ""]

model = pull_from_hub(repo_id="Dimitre/mobilenet_v3_small")

def preprocess(image):
    image = image.reshape((-1, 224, 224, 3)) # (batch_size, height, width, num_channels)
    return image / 255.

def postprocess(prediction):
    return {labels[i]: prediction[i] for i in range(len(labels))}
    # return {labels[i]: 0 for i in range(len(labels))}

def predict_fn(image):
    image = preprocess(image)
    prediction = model(image)[0].numpy()
    print('****************')
    print(prediction)
    try:
        print("default")
        print({labels[i]: prediction[i] for i in range(len(labels))})
    except:
        print("default gives error")
    print('****************')
    print(list(prediction))
    try:
        print("list")
        print({labels[i]: list(prediction)[i] for i in range(len(labels))})
    except:
        print("list gives error")
    scores = postprocess(prediction)
    return scores

iface = gr.Interface(fn=predict_fn, 
                     inputs=gr.Image(shape=(224, 224)), 
                     outputs=gr.Label(num_top_classes=5))
iface.launch()