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


__all__ = ['load_learner', 'lean', 'greet', 'classify_image',
           'image', 'label', 'examples', 'interface']

path = Path(__file__).parent


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


learn_inf = load_learner(path/'plane.pkl')
# learn_inf.predict('test.png')


# print(learn_inf.predict('test.png'))
# iface = gr.Interface(fn=greet, inputs="text", outputs="text")
# iface.launch()

# categories = ['F-22 Raptor', 'SR-71 Blackbird', 'B-2 Spirit']
categories = learn_inf.dls.vocab


def classify_image(img):
    pred, pred_idx, probs = learn_inf.predict(img)
    return {categories[i]: float(probs[i]) for i in range(len(probs))}


image = gr.inputs.Image(shape=(224, 224))
label = gr.outputs.Label(num_top_classes=3)


examples = [['b2.jpg'], ['f22.jpg'], ['sr71.jpg']]

interface = gr.Interface(fn=classify_image, inputs=image,
                         outputs=label, examples=examples)
interface.launch(inline=False)


# print(learn_inf.predict('b2.jpg'))