## importing dependencies
import gradio as gr
from fastai.vision.all import *

## importing model created
learn = load_learner('model.pkl')

## creating labels and examples to be used by the interface
labels = learn.dls.vocab
examples = ['stop_sign.jpg',
            'pedestrian_crossing_sign.jpg',
            'speed_limit_sign.jpg',
            'bump_sign.jpg',
            'semaphore_sign.jpg']

## creating prediction function for the interface
def predict(img):
    img = PILImage.create(img)
    pred, pred_idx, probability = learn.predict(img)
    return {labels[i]: float(probability[i]) for i in range(len(labels))}

## generating huggingFace's interface
gr.Interface(fn=predict,
             inputs=gr.components.Image(shape=(512, 512)),
             outputs=gr.components.Label(num_top_classes=3),
             examples=examples
             ).launch(share=False)