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

learn = load_learner("./export.pkl")

categories = ("Black", "Grizzly", "Teddy")


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


image = gr.inputs.Image(shape=(192, 192))
label = gr.outputs.Label()

iface = gr.Interface(
    fn=classify_image,
    inputs=image,
    outputs=label,
    title="Bear Classifier",
    description="Upload a photo of either a Black bear, Grizzly bear, or Teddy bear, and submit to have the model detect the bear.",
)
iface.launch()