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

learn = load_learner("model.pkl")

def is_cat(x):
    return x[0].isupper()

labels = learn.dls.vocab

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

iface = gr.Interface(
    fn=predict,
    inputs=gr.Image(shape=(512, 512)),
    outputs=gr.Label(num_top_classes=2, label="Is it a cat?"),
    title="Is it a cat or a dog?!",
    description="An app to classify whether it is a cat or a dog. The model is trained on the Oxford Pets dataset with fastai.",
    examples=["download.jpg", "33669.jpg"],
    interpretation="default",
)
iface.launch(enable_queue=True)