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

learn = load_learner("model.pkl")


def predict(img):
    # use PIL to resize the image to 224x224
    img = Image.fromarray(img)
    pred, pred_idx, probs = learn.predict(img)
    # if the prediction contains poodle, return poodle
    if "poodle" in pred:
        return "poodle"
    return "doodle"


title = "Poodle or Doodle"
description = "A resnet18 model fine tuned on 400 images of poodles and labradoodles. You can predict using the default image or upload an image an image of your own to see if you're looking at a poodle or doodle."

gr.Interface(
    fn=predict,
    inputs=gr.Image(value="labradoodle.jpg"),
    outputs=gr.Label(),
    title=title,
    description=description,
    allow_flagging="never",
).launch()