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

learn = load_learner("model.pkl")

categories = learn.dls.vocab


def predict(img):
    pred, pred_idx, probs = learn.predict(img)
    return dict(zip(categories, map(float, probs)))


title = "Bear Classifier"
description = "A bear classifier trained on the images downloaded from bing serach with fastai."

image = gr.inputs.Image(shape=(192, 192))
label = gr.outputs.Label()
examples = ["black_bear.jpg", "brizzly_bear.jpg", "teddy_bear.jpg"]
interpretation = "default"
enable_queue = True

itf = gr.Interface(fn=predict,
                   inputs=image,
                   outputs=label,
                   title=title,
                   examples=examples,
                   interpretation=interpretation,
                   enable_queue=enable_queue)

itf.launch()