from pathlib import Path
from fastai.vision.learner import load_learner
import gradio as gr

MODELPATH='cygnet-vs-duckling.pkl'
EXAMPLEFOLDER = Path().cwd()

learn = load_learner(MODELPATH)
categories = learn.dls.vocab

def classify_image(image):
    _, _, probs = learn.predict(image)
    return dict(zip(categories, map(float, probs)))

title = 'Mirror, mirror on the wall am I a duckling or a cygnet after all?'
description = """Hans Christian Andersen's tale of the ugly duckling tells us about the sad youth of a cygnet acidentally being brought up in a family of ducks and being ostrized for its ugliness.\n But what if the poor cygnet had a magic mirror to tell it that it was in fact a young swan after all? Machine learning to the rescue!"""

examples = EXAMPLEFOLDER.glob('*.jpg')

app = gr.Interface(fn=classify_image,
                   inputs=gr.components.Image(),
                   outputs=gr.components.Label(),
                   examples=examples,
                   title=title,
                   description=description,
                   allow_flagging='never')

app.launch()