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

# returns dict of predictions and their probabilities sorted
def classify_img(img):
    instrument,index,probs = learn.predict(PILImage.create(img))    
    sorted_probs = sorted(zip(learn.dls.vocab, map(float,probs)), key=lambda x: x[1], reverse=True)
    return dict(sorted_probs)
    
#paths
mdl = 'minst.pk1'
samples = 'examples/'

# load pre-trained model 
learn = load_learner(mdl)

examples = [f'{samples}acoustic guitar.jpg', f'{samples}drums.jpg', f'{samples}electric guitar.jpg', f'{samples}keyboard.jpg']

demo = gr.Interface(fn=classify_img, inputs=gr.Image(width=224,height=224), outputs="label", examples=examples, title="Musical Instrument Classifier")
demo.launch()