import gradio as gr
#from transformers import pipeline
from tensorflow.keras.models import load_model

#pipe = pipeline(task="image-classification", model="SuperSecureHuman/Flower-CNN")

model = load_model('./model.h5')


def predict_image(img):
    img_4d = img.reshape(-1, 224, 224, 3)
    prediction = model.predict(img_4d)[0]
    return {class_names[i]: float(prediction[i]) for i in range(10)}


class_names = ['Crescentia_Cujete', 'Fiddle_Wood', 'Gold_Apple', 'Hill_Mango', 'Indian_Tulip_Tree', 'Mahagony', 'Pala_Indigo_Plant', 'Spanish_Cherry', 'Teak', 'Yellow_Trumpet']

image = gr.inputs.Image(shape=(224, 224))

label = gr.outputs.Label(num_top_classes=10)

gr.Interface(fn=predict_image,
             title="Tree Classification",
             description="Tree CNN",
             inputs=image,
             outputs=label,
             live=True,
             interpretation='default',
             allow_flagging="never").launch()