import os
import gradio as gr
import numpy as np
import tensorflow as tf
from tensorflow.keras import models

IMG_SIZE = 300
class_names = ['none','mild','severe']

cwd = os.getcwd()
outpath= os.path.join(cwd,"model")
model_name = 'cross_event_ecuador_efficientnet_1643984852.h5'
loaded_model = models.load_model(os.path.join(outpath,model_name))

def _classifier(inp):
  img = np.asarray(tf.cast(inp, dtype=tf.float32)) * 1 / 255.0
  img = img.reshape((-1, IMG_SIZE, IMG_SIZE, 3))
  preds = loaded_model.predict(img).flatten()
  return {class_names[i]:float(preds[i]) for i in range(len(class_names))}

iface = gr.Interface(fn=_classifier, 
                     inputs=gr.inputs.Image(shape=(IMG_SIZE, IMG_SIZE)),
                     outputs=gr.outputs.Label()
                     )

iface.launch()