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 that tells if something's huggable or not?",
    article = "<p style='text-align: center'><a href='https://github.com/daspartho/is-it-huggable' target='_blank'>Github</a></p>",
    fn=classify_image, 
    inputs=gr.Image(shape=(224,224)), 
    outputs=gr.outputs.Label(),
    examples=['examples/dog.jpg', 'examples/cactus.jpg', 'examples/plushie.jpg', 'examples/snowman.jpg', 'examples/chainsaw.jpg', 'examples/shark.jpg','examples/bunny.jpg', 'examples/knife.jpg', 'examples/tiger.jpg', 'examples/trex.jpg'],
    live=True,
    enable_queue=True
    )
iface.launch()