from fastai.vision.all import*
import gradio as gr
learn1 = load_learner('stage1.pkl')

categories1 = 'discarded clothing', 'food waste', 'plastic bags', 'recyc_no_scrap', 'scrap metal piece', 'wood scraps'


def classify_stage1(img):
    pred, idx, probs = learn1.predict(img)
    return dict(zip(categories1, map(float,probs)))

image1 = gr.inputs.Image(shape=(192,192))
label1 = gr.outputs.Label()
examples1 = ['stage1ex1_t.jpeg', 'stage1ex2_t.jpeg','stage1ex3_t.jpeg','stage1ex4_t.jpeg', 'stage1ex5_t.jpeg','stage1ex6_t.jpeg']
intf1 = gr.Interface(fn=classify_stage1, inputs = image1, outputs = label1, examples = examples1)
intf1.launch(inline=False)