import time
import gradio as gr
import tensorflow as tf
import matplotlib.pyplot as plt

def predict(input_img):

    start = time.time()
    # Load the saved Keras model
    model = tf.keras.models.load_model("VGG19.h5")
    # Preprocess the image
    # img_0 = tf.keras.utils.load_img(input_image)
    img_0 = tf.keras.utils.img_to_array(input_img)
    img_0 = tf.image.resize(img_0, (256, 256))
    img_1 = tf.expand_dims(img_0, axis = 0)

    class_names = ["bordered", "borderless", "row_bordered"]

    # Make predictions using the model
    predictions = model.predict(img_1)
    predicted_label = tf.argmax(predictions, 1).numpy().item()

    for item in predictions :
      item = tf.round((item*100))
    
    fig = plt.figure(1, figsize=(8, 14))
    plt.axis('off')
    plt.rcParams.update({'font.size': 24})
    plt.title(f'prediction : {class_names[predicted_label]}\n\n'
              f'{item[0]} % {class_names[0]}\n'
              f'{item[1]} % {class_names[1]}\n'
              f'{item[2]} % {class_names[2]}\n')
    plt.imshow(img_0/255)
    end = time.time()
    inf_time = end - start
    return plt, inf_time

    
output = gr.Plot()
gradio_app = gr.Interface(
    predict,
    inputs=gr.Image(label="table type", sources=['upload', 'webcam'], type="pil"),
    outputs=[output, gr.Textbox(label="Inference time")],
    title="Table-type Classification"
)

if __name__ == "__main__":
    gradio_app.launch(share=True, debug=True)