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


learn = load_learner('kitten.pkl')
classes = learn.dls.vocab

def classify_images(img):
    img = PILImage.create(img)
    pred,idx,prob = learn.predict(img)
    return {classes[i]: float(prob[i]) for i in range(len(classes))}

title = "Cute or Ugly Kitten Classifier"
description = "Upload a kitten and it will tell you if it's ugly or cute! ~Elio."
examples = [['cute kitten.jpg'], ['ugly kitten.jpg']]

iface = gr.Interface(fn=classify_images, inputs=gr.inputs.Image(shape=(512, 512)), outputs=gr.outputs.Label(), title=title, description=description, examples=examples)
iface.launch()