import os

import gradio as gr
import numpy as np
import pandas as pd
import tensorflow as tf
from tensorflow.keras.layers import *
from tensorflow.keras.preprocessing.image import ImageDataGenerator


def classify_grapevine_leaves(img):
    categories = ("Healthy", "Powdery Mildew", "Rust")

    # load keras model
    model = tf.keras.models.load_model("./keras_model/")
    # load image
    # img = tf.keras.preprocessing.image.load_img(img, target_size=(360, 360))
    # convert image to array
    img = tf.keras.preprocessing.image.img_to_array(img)
    # add batch dimension
    img = tf.expand_dims(img, axis=0)
    # predict
    prediction = model.predict(img)
    # get label
    print(np.argmax(prediction, axis=1))
    label = categories[prediction.argmax()]
    # get confidence
    conf = prediction[0][prediction.argmax()]
    # return label and confidence
    return dict(zip(categories, map(float, prediction[0])))


exemples = [
    "./exemples/healthy.jpg",
    "./exemples/powdery.jpg",
    "./exemples/rust.jpg",
]

image = gr.inputs.Image(shape=(360, 360))
label = gr.outputs.Label()
app = gr.Interface(
    fn=classify_grapevine_leaves, inputs=image, outputs=label, examples=exemples
)
app.launch(inline=False)