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', 'sunflower.jpg', 'whiteclouds.jpg', 'yellowclouds.jpg']

article = 'This model was build using a resnet-18 architecture with weights pretrained on the ImageNet data set and fine-tuned using about 80 images of ducklings and cygnets each.\nNote that it is binary classifier and can therefore only output cygnet or duckling, "other" is not an option. As a fun exercise, I included some non-waterfowl pictures in the example. Can you guess what the model will classify them as?\nOn a final note, whilst this classifier claims to be able to detect ducklings, it really only detects mallard ducklings (aka the yellow ones) and has a hard time recognizing ducklings of other species. To see this in action, compare its performance for a mallard duckling with its performance when given the image of a black cayuga duckling for example.'

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