test

app.py

@@ -1,24 +1,18 @@
-__all__ = ['models', 'image', 'label', 'examples', 'intf']
-
 import gradio as gr
 from fastai.vision.all import load_learner
 
-models = [
-    load_learner('model-color.pkl'),
-    load_learner('bricks-model.pkl')
-]
+def classify_image_color(img):
+    learn_color = load_learner('lego-bricks-detecter-model-color.pkl')
+    pred,_,probs = learn_color.predict(img)
+    return dict(zip(learn_color.dls.vocab, map(float,probs)))
 
-def classify_image(img):
-    results = {}
-    for i, learn in enumerate(models):
-        categories = learn.dls.vocab
-        pred, idx, probs = learn.predict(img)
-        results[f'model{i+1}'] = dict(zip(categories, map(float,probs)))
-    return results
+def classify_image_shape(img):
+    learn_shape = load_learner('lego-bricks-detecter-model-shape.pkl')
+    pred,_,probs = learn_shape.predict(img)
+    return dict(zip(learn_shape.dls.vocab, map(float,probs)))
 
 image = gr.inputs.Image(shape=(256, 256))
-label = gr.outputs.Label()
-examples = []
+label = gr.outputs.Label(num_top_classes=3)
 
-intf = gr.Interface(fn=classify_image, inputs=image, outputs=label, examples=examples)
-intf.launch(inline=False)
+intf = gr.Interface(fn=[classify_image_color, classify_image_shape], inputs=image, outputs=label, examples="", title="Lego Brick Classifier", layout="vertical")
+intf.launch(share=True)