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

categories2 = 'HDPE container', 'PET plastic bottle', 'aluminium can', 'cardboard', 'glass', 'paper2D', 'paper3D', 'steel and tin cans'


def classify_stage2(img):
    pred, idx, probs = learn2.predict(img)
    return dict(zip(categories2, map(float,probs)))

image2 = gr.inputs.Image(shape=(192,192))
label2 = gr.outputs.Label()
examples2 = ['stage2ex1_t.jpeg', 'stage2ex2_t.jpeg','stage2ex3_t.jpeg', 'stage2ex4_t.jpeg','stage2ex5_t.jpeg',
            'stage2ex6_tt.jpeg','stage2ex7_tt.jpeg','stage2ex8_t.jpeg']
intf2 = gr.Interface(fn=classify_stage2, inputs = image2, outputs = label2, examples = examples2)
intf2.launch(inline=False)