import gradio as gr
from fastai.vision.learner import load_learner

model = load_learner("model.pkl")


def classify_image(image):
    preds = model.predict(image)[2]
    return dict(zip(["Bird", "Forest"], preds.tolist()))


gr.Interface(
    fn=classify_image,
    inputs=gr.Image(),
    outputs=gr.Label(num_top_classes=2),
    examples=["images/bird.jpeg", "images/forest.jpeg"],
    title="Bird or Forest?",
).launch()