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

learner_inf = load_learner('model.pkl')

labels = learner_inf.dls.vocab

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

title = "Cat vs. Dog Classifier"
description = "A Cat/Dog Classfier trained on random images obtained from DuckDuckGo. Used as a demo for Gradio + HuggingFace."
interpretation='default'

gr.Interface(
  fn=predict, inputs=gr.inputs.Image(shape=(512, 512)),
  outputs=gr.outputs.Label(num_top_classes=2),
  title=title,
  description=description,
  interpretation=interpretation
).launch()