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

learn = load_learner('BF_vs_MTH_model.pkl')

labels = learn.dls.vocab

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

title = 'Butterflies vs Moths'
description = 'A basic app which checks whether the image uploaded by you is of a butterfly or a moth!'
examples = ['Butterfly_1.png', 'Butterfly_2.png', 'Moth_1.png', 'Moth_2.jpg']

UI = gr.Interface(
    fn=predict,
    inputs=gr.inputs.Image(shape=(512,512)),
    outputs=gr.outputs.Label(num_top_classes=3),
    title=title,
    description=description,
    examples=examples
)

UI.launch()