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

learn_inf = load_learner('model.pkl')

features = learn_inf.dls.vocab

def predict_image(image):
    pred_class, pred_idx, probs = learn_inf.predict(image)
    return {str(features[i]): float(probs[i]) for i in range(len(learn_inf.dls.vocab))}

input_image = gr.inputs.Image(shape=(None, None))
output_label = gr.outputs.Label(num_top_classes=3)

gr.Interface(fn=predict_image,
             inputs=input_image,
             outputs=output_label,
             title="Image Classifier",
             description=f"Upload an image to classify it into one of the following categories: {' | '.join([f'{item}' for item in features])}",
             capture_session=True).launch()