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

# import model for gradio
learn_gradio = load_learner('image_classifier_flowers.pkl')

# build prediction function
labels = learn_gradio.dls.vocab

def predict(img):
    img = PILImage.create(img)
    pred,pred_idx,probs = learn_gradio.predict(img)
    return {str(labels[i]): float(probs[i]) for i in range(len(labels))}

# build gradio interface
gradio_interface = gr.Interface(
    title = "Flower Image Classifier",
    description = "A simple classifier for the 102-category <a href='https://www.robots.ox.ac.uk/~vgg/data/flowers/' target='new'>Flower Dataset</a>",
    fn=predict, 
    inputs = gr.inputs.Image(shape=(224,224)), 
    outputs = gr.outputs.Label(num_top_classes=5)
)

# launch interface
gradio_interface.launch(enable_queue=True)