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

MODELPATH='cygnet-vs-duckling.pkl'

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 which is accidentally brought up in a family of ducks and is ostrized on the account of it being different. But what if the cygnet had a magic mirror to tell it that it was in fact a young swan after all? Machine learning to the rescue!"""

examples = ['duckling.jpg', 'cygnet.jpg', 'sunflower.jpg', 'whiteclouds.jpg', 'yellowclouds.jpg']

article = 'This model was build using a resnet-18 architecture, pretrained on ImageNet and finetuned using 100 images of ducklings and cygnets each.'

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

app.launch()