import gradio as gr
import os

from fastai.learner import load_learner
from pathlib import Path

path = Path(os.path.dirname(os.path.realpath(__file__)))

# Make sure to export the model first via the notebook.
# The exported model is not included in the repository.
model = load_learner(path/'export.pkl')
labels = model.dls.vocab

def predict(img):
    _, _, probabilities = model.predict(img)
    return {labels[i]: float(probabilities[i]) for i in range(len(labels))}

title = "Grizzly, Black, or Teddy Bear?"
description = "This is a simple image classifier that can predict whether a given image contains a grizzly bear, black bear, or teddy bear."

gr.Interface(
    fn = predict,
    inputs = gr.Image(width = 512, height = 512),
    outputs = gr.Label(num_top_classes = 3),
    title = title,
    description = description,
    allow_flagging = 'never'
).launch()