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

learn = load_learner('model.pkl')

labels = learn.dls.vocab
def predict(img):
    img = PILImage.create(img)
    pred,pred_idx,probs = learn.predict(img)
    return {labels[i]: float(probs[i]) for i in range(len(labels))}
  
  
with gr.Blocks(theme=gr.themes.Monochrome()) as demo:
    img = gr.Image(height=500, sources=["upload"])
    examples = gr.Examples([
        ["./imgs/dali2.jpeg"],
        ["./imgs/dali3.jpeg"],
        ["./imgs/dali4.jpeg"],
        ["./imgs/dali5.jpeg"],
        ["./imgs/dali6.jpeg"],
    ], inputs=img)
    with gr.Row():
        clear_btn = gr.ClearButton(value="Clear")
        submit_btn = gr.Button("Predict")
    output = gr.Label(num_top_classes=2, label="Labels")
    submit_btn.click(predict, img, output)
    clear_btn.add(img)

demo.launch()