import tensorflow as tf
import gradio as gr
import cv2
import os

used_model = tf.keras.layers.TFSMLayer(os.path.dirname('/model'), call_endpoint='serving_default')
new_classes = ['blight', 'common_rust', 'gray_leaf_spot','healthy']

def classify_image(img_dt):
  img_dt = cv2.resize(img_dt,(256,256))
  img_dt = img_dt.reshape((-1,256,256,3)) 
  prediction = used_model.predict(img_dt).flatten()
  confidences = {new_classes[i]: float(prediction[i]) for i in range (4) }
  return confidences


with gr.Blocks() as demo:
    with gr.Row():
        signal = gr.Markdown(''' #Welcome to Maize Classifier, This model can identify if a leaf is 
        **HEALTHY**, has **COMMON RUST**, **BLIGHT** or **GRAY LEAF SPOT**''')
        inp = gr.image()
        out = gr.Label()
        inp.upload(fn= classify_image, inputs = inp, outputs = out, show_progrss = True)