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

def is_cat(x): return x[0].isupper() # Required for model
    
learn = load_learner('model.pkl')

# Use labels created by me for predictions
labels = learn.dls.vocab
label_dict = {0:'PERRO',1:'GATO'}

def predict(img):
	img = PILImage.create(img).resize((512, 512))
	pred, pred_idx, probs = learn.predict(img)
	return {label_dict.get(labels[i],'DUNNO'): float(probs[i]) for i in range(len(labels))}

title = "Pet Breed Classifier"
description = "A pet breed classifier trained on the Oxford Pets dataset with fastai."

gr.Interface(
fn=predict,
inputs=gr.Image(),
outputs=gr.Label(num_top_classes=3),
title=title,
description = description).launch(share=True)