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

title = "Harry Potter Character classifier"
description = "A Harry Potter Character classifier trained with over 3000 images from the internet!"
article = "<p style='text-align: center'><a href='https://aldrinjenson.me' target='_blank'>Created by Aldrin Jenson</a></p>"
examples = ['sample_images/daniel_radcliffe-000360.jpeg', 'sample_images/daniel_radcliffe-000003.jpeg', 'sample_images/daniel_radcliffe-000362.jpeg', 'sample_images/lunaLovegood-000003.jpeg', 'sample_images/groupPhoto-000426.jpeg', 'sample_images/daniel_radcliffe-000504.jpeg', 'sample_images/dracoMalfoy-000002.jpeg', 'sample_images/ginnyweasly-000365.jpeg', 'sample_images/daniel_radcliffe-000324.jpeg', 'sample_images/ginnyweasly-000328.jpeg', 'sample_images/hermionie-000001.jpeg', 'sample_images/ginnyweasly-000003.jpeg', 'sample_images/ginnyweasly-000404.jpeg', 'sample_images/lunaLovegood-000001.jpeg', 'sample_images/lunaLovegood-000002.jpeg',
            'sample_images/dracoMalfoy-000003.jpeg', 'sample_images/daniel_radcliffe-000001.jpeg', 'sample_images/groupPhoto-000423.jpeg', 'sample_images/daniel_radcliffe-000004.jpeg', 'sample_images/daniel_radcliffe-000466.jpeg', 'sample_images/ginnyweasly-000002.jpeg', 'sample_images/daniel_radcliffe-000202.jpeg', 'sample_images/groupPhoto-000515.jpeg', 'sample_images/groupPhoto-000464.jpeg', 'sample_images/ginnyweasly-000317.jpeg', 'sample_images/ginnyweasly-000001.jpeg', 'sample_images/daniel_radcliffe-000381.jpeg', 'sample_images/ginnyweasly-000139.jpeg', 'sample_images/ginnyweasly-000388.jpeg', 'sample_images/daniel_radcliffe-000002.jpeg']
interpretation = 'default'
enable_queue = True


learn = load_learner('export.pkl')
labels = learn.dls.vocab


def predict(img):
    img = PILImage.create(img)
    pred, pred_idx, probs = learn.predict(img)
    return {labels[i]: float(probs[i]) for i in range(len(labels))}


interface = gr.Interface(fn=predict, inputs=gr.components.Image(shape=(512, 512)), outputs=gr.components.Label(
    num_top_classes=4), title=title, description=description, article=article, examples=examples, interpretation=interpretation)
interface.launch(share=True, enable_queue=enable_queue)