import tensorflow.compat.v2 as tf
import tensorflow_hub as hub
import numpy as np
import pandas as pd
import cv2
import gradio as gr

model = hub.KerasLayer('https://tfhub.dev/google/aiy/vision/classifier/food_V1/1')

labelmap_url = "https://www.gstatic.com/aihub/tfhub/labelmaps/aiy_food_V1_labelmap.csv"
input_shape = (224, 224)
classes = list(pd.read_csv(labelmap_url)["name"])

def classify_image(image):
    image = image.reshape((-1, 224, 224, 3))
    image = image / image.max()
    output = model(image)
    output = list(output)[0]
    return {classes[i]: float(output[i]) for i in range(len(classes))}
    
image = gr.inputs.Image(shape=(224, 224))
label = gr.outputs.Label(num_top_classes=8)

gr.Interface(
    fn=classify_image,
    inputs=image,
    outputs=label,
    description="Demo of TF Hub-hosted model aiy/vision/classifier/food_V1",
    article="This demo is a proof-of-concept inspired in https://tfhub.dev/google/aiy/vision/classifier/food_V1/1",
    examples=[["cake.jpg"]],
).launch()