from fastai.vision.all import *
import gradio as gr

learn = load_learner('model_convnext_small_in22k_version_1.pkl')
categories = ("fire", "nofire")

def classify_image(img):
    pred, idx, prob = learn.predict(img)
    return dict(zip(categories, map(float, prob)))

img = gr.inputs.Image(shape=(250, 250))
label = gr.outputs.Label()

gr.Interface(
    fn=classify_image, 
    inputs="image", 
    outputs="label",
    enable_queue=True,
    examples=["after-a-bushfire.jpg", "fire_1.jpg", "forest_1.jpg", "sebastian-unrau-sp-p7uuT0tw-unsplash.jpg"],
    interpretation='default').launch()