import gradio as gr
from fastai.vision.all import load_learner, get_image_files

learn = load_learner('resnet18_bearclassifier_linux.pkl')

# gradio UI


def classify_image(img):
    categories = ['brown', 'polar']
    _, _, probability = learn.predict(img)
    return dict(zip(categories, map(float, probability)))


app = gr.Interface(fn=classify_image, inputs='image',
                   outputs='label', examples=list(get_image_files('valid')))
app.launch()
classify_image('images.jpg')