app.py

__all__ = ['learn', 'classify_image', 'categories', 'image', 'label', 'examples', 'intf']

import gradio as gr

def classify_image(img):
    from fastai.vision.all import load_learner
    learn = load_learner('lego-bricks-model.pkl')
    categories = (
     #'11214 Bush 3M friction with Cross axle',
     #'18651 Cross Axle 2M with Snap friction',
     #'2357 Brick corner 1x2x2',
     #'3003 Brick 2x2',
     #'3004 Brick 1x2', 
     #'3005 Brick 1x1',
     #'3022 Plate 2x2',
     #'3023 Plate 1x2',
     #'3024 Plate 1x1',
     #'3040 Roof Tile 1x2x45deg',
     #'3069 Flat Tile 1x2',
     #'32123 half Bush',
     #'3673 Peg 2M',
     #'3713 Bush for Cross Axle',
     #'3794 Plate 1X2 with 1 Knob',
     #'6632 Technic Lever 3M', 
    'Brick_502','Brick_602','Brick_685','Brick_711','Brick_938','Brick_2499','Brick_2610','Brick_2921','Brick_3479','Brick_3708','Brick_4503',
    'Brick_4515','Brick_6162','Brick_6182','Brick_6191','Brick_11476','Brick_12897','Brick_12899','Brick_13760','Brick_16615','Brick_18896','Brick_18976',
    'Brick_21445','Brick_21980','Brick_24869','Brick_26280','Brick_28324','Brick_30357','Brick_30407','Brick_31520','Brick_32124','Brick_41630','Brick_41748',
    'Brick_42022','Brick_42074','Brick_42604','Brick_45407','Brick_45590','Brick_48002','Brick_50956','Brick_52216','Brick_54671','Brick_81599','Brick_87995',
    'Brick_88704','Brick_92092','Brick_94318','Brick_303226','Brick_4106592','Brick_4261453'
) 
    pred,idx,probs = learn.predict(img)
    return dict(zip(categories, map(float,probs)))

image = gr.inputs.Image(shape=(192, 192))
label = gr.outputs.Label()
examples = ['0002.png', '201706161606-0002.png', '201706161906-0001.png']

intf = gr.Interface(fn=classify_image, inputs=image, outputs=label, examples=examples)
intf.launch(inline=False)