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

learn = load_learner('archi_styles_model.pkl')

styles = ('Baroque Architecture','Gothic Architecture', 'Greek and Roman Classical Architecture', 'Modern Architecture', 'Neoclassical Architecture', 'Neofuturist Architecture', 'Post-Modern Architecture', 'Victorian Architecture')
examples = ['Baroque Architecture.jpg', 'Gothic Architecture.jpg', 'Modern Architecture.jpg']

# Use dict() and zip() to build a dict, not formaliar with data object in python.
def image_classifier(img):
    pred_styles,pred_idx,probs = learn.predict(img)
    return dict(zip(styles, map(float, probs)))

demo = gr.Interface(fn=image_classifier, inputs="image", outputs="label", examples = examples)
demo.launch()