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 had a magic mirror to tell it that it had been a young swan all along? Machine learning to the rescue!"""

examples = ['duckling.jpg', 'cygnet.jpg']

article = '**Technical Details**: The classification model was build using a resnet-18 architecture. Training was done using a transfer learning approach. I took the publicly available weights that were pre-trained on the ImageNet data set and fine-tuned them using about 80 images of ducklings and cygnets each.\nNote that it is binary classifier and will therefore only output "cygnet" or "duckling", "other" is not an option.' 

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()