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



learn = load_learner('private.pkl')
categories = learn.dls.vocab
def classify_image(img):
  pred, idx, probs = learn.predict(img)
  #gradio will return a dictionary, and change tensors from pytorch to float
  return dict(zip(categories, map(float, probs)))


demo = gr.Interface(
    fn=classify_image,
    inputs="image",
    outputs="label"
)
demo.launch(share=True)