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

learn_inf = load_learner('model.pkl')
categories = learn_inf.dls.vocab

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

iface = gr.Interface(
    title = "Is it Huggable?",
    description = "An image classifier to classify things as huggable or not, trained on examples of both categories.",
    fn=classify_image, 
    inputs=gr.inputs.Image(shape=(224,224)), 
    outputs=gr.outputs.Label(),
    examples=['lego.jpg', 'dog.jpg', 'cactus.jpg', 'plushie.jpg', 'snowman.jpg'],
    live=True,
    enable_queue=True
    )
iface.launch()