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

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

article="<p style='text-align: center'><a href='https://www.linkedin.com/in/zakcroft' target='_blank'>Me</a></p>"

demo = gr.Interface(fn=predict,
                    inputs=gr.Image(width=512, height=512),
                    outputs=gr.Label(num_top_classes=3),
                    title="Bear Classifier",
                    description="Grizzly, black or Teddy?",
                    article=article,
                    examples=['black.jpeg', 'grizzly.jpeg', 'teddy.jpeg',
                              'black-2.jpeg', 'grizzly-2.webp', 'teddy-2.webp'],
                    # interpretation='default',
                    # enable_queue=True
                    flagging_options=['should be Black', 'should be Grizzly', 'should be Teddy', 'should be None']
                    ).launch(share=True)


# demo = gr.Interface(
#     fn=predict,
#     inputs=gr.inputs.Image(shape=(512, 512)),
#     outputs=gr.outputs.Label(num_top_classes=3),
#     title=title,
#     description=description
#     ,article=article,
#     examples=examples,
#     interpretation=interpretation,
#     enable_queue=enable_queue).launch()

if __name__ == "__main__":
    demo.launch(show_api=False)