from fastai.vision.all import *

learn = load_learner('export.pkl')

# searches = 'steve minecraft','herobrine minecraft'

def sanity_check():
    for o in searches:
        category,_,probs = learn.predict(PILImage.create(f'{o}.jpg'))
        print(f"This is a: {category}.")
        print(f"Probability it's a {category}: {probs[0]:.4f}")

# sanity_check()

labels = learn.dls.vocab

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

import os
os.system("pip install --upgrade --force-reinstall gradio==3.50")

import gradio as gr

gr.Interface(
    fn=predict, 
    inputs=gr.inputs.Image(shape=(512, 512)), 
    examples = [ 'steve minecraft.jpg', 'herobrine minecraft.jpg' ],
    interpertation='default',
    outputs=gr.outputs.Label(num_top_classes=3)
).launch(share=True)