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

# gradio
learn_caltech101 = load_learner('image_classifier_caltech101.pkl')

# build prediction function
labels = learn_caltech101.dls.vocab

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

# build gradio interface
gradio_interface = gr.Interface(
    title = "Caltech_101 Image Classifier",
    description = "A simple image classifier based on the Caltech_101 dataset.",
    fn=predict, 
    inputs = gr.inputs.Image(shape=(224,224)), 
    outputs = gr.outputs.Label(num_top_classes=5)
)

gradio_interface.launch(enable_queue=True)