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


def greet(name):
    return "Hello " + name + "!!"


learn = load_learner("model-v7.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))}


examples = [
    "a2_b_m.jpg",
    "a2_coeur_lyon.jpg",
    "gzup_america.jpg",
    "gzup_gameboy.jpg",
    "pa_425.jpg",
    "pa_1341.jpg",
    "stork_music.jpg",
    "stork_prince.jpg",
]

DESCRIPTION = """
Street art deep learning model. It's trained to recognize the style of the following artists active in Paris:
[A2](https://a2-streetart.com/), [gzup](https://www.gzup.fr/in-the-streets/), [invader](https://www.instagram.com/invaderwashere), [stork](https://www.instagram.com/stork_pixelart/).
For more AI experiments check out my [newsletter](https://newsletter.pnote.eu) 💫.
"""

gr.Interface(
    fn=predict,
    inputs=gr.inputs.Image(shape=(512, 512)),
    outputs=gr.outputs.Label(num_top_classes=3),
    examples=examples,
    title="Art reco, recognizer of street art",
    description=DESCRIPTION
).launch(share=True)