import gradio as gr
import tensorflow as tf
from tensorflow import keras

import huggingface_hub
from huggingface_hub import from_pretrained_keras


# load custom pre-trained model from HuggingFace models
model_api_link = 'chaninder/waste-sorting-model-v4' #'chaninder/waste-sorting-model-updated', #'chaninder/waste-sorting-model'
pre_trained_model = from_pretrained_keras(model_api_link)

# classification labels
labels = ['compost', 'e-waste', 'recycle', 'trash']


def classify_image(inp):
    inp = inp.reshape((-1, 224, 224, 3))
    #inp = tf.keras.applications.mobilenet_v2.preprocess_input(inp)
    prediction = pre_trained_model.predict(inp).flatten()
    confidences = {labels[i]: float(prediction[i]) for i in range(4)}
    return confidences

# create Gradio interface

iface = gr.Interface(fn=classify_image, 
             inputs=gr.Image(shape=(224, 224)),
             outputs=gr.Label(num_top_classes=4),
             examples=["banana.jpg", 'can.jpg', 'battery.jpg'])
             
iface.launch(share=True)