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


learn = load_learner("bear-classifier.pkl")
labels = learn.dls.vocab

def predict(img):
    """
    Prediction APIs.
    """
    _, _, probs = learn.predict(PILImage.create(img))

    return {labels[i]: float(probs[i]) for i in range(len(labels))}


demo = gr.Interface(
    fn=predict, 
    inputs=gr.Image(), 
    outputs=gr.Label(num_top_classes=3), 
    title="Bear Classifier", 
    description="A bear classifier fine tuned on ResNet18 with a few bear samples from the internet. Its task is to recognize whether an image is a grizzy bear, a black bear or just a cute litte teddy bear!", 
    examples=["grizzly.jpg"]
)
demo.launch()