import gradio as gr
from tensorflow import keras
import tensorflow as tf

model = tf.keras.models.load_model("./model.h5")

#return a dictionary of labels and probabilities
def fire_or_non_fire(img):
    img = img.reshape((-1 ,105, 165, 3))
    prediction = model.predict(img).tolist()[0]
    labels = ["Fire", "Non Fire"]
    return {labels[i]: prediction[i] for i in range(2)}

im = gr.inputs.Image(shape=(105, 165), source="upload")
iface = gr.Interface(
    fn = fire_or_non_fire, 
    inputs = im, 
    outputs = gr.outputs.Label(),
)

iface.launch()