asdf

app.py

@@ -1,18 +1,41 @@
 import gradio as gr
+from gradio.mix import Parallel
 from fastai.vision.all import load_learner
 
 def classify_image_color(img):
-    learn_color = load_learner('model-color.pkl')
-    pred,_,probs = learn_color.predict(img)
-    return dict(zip(learn_color.dls.vocab, map(float,probs)))
+    from fastai.vision.all import load_learner
+    learn = load_learner('model-color')
+    categories = learn.dls.vocab
+    pred, idx, probs = learn.predict(img)
+    return {category: float(prob) for category, prob in zip(categories, probs)}
 
 def classify_image_shape(img):
-    learn_shape = load_learner('bricks-model.pkl')
-    pred,_,probs = learn_shape.predict(img)
-    return dict(zip(learn_shape.dls.vocab, map(float,probs)))
+    from fastai.vision.all import load_learner
+    learn = load_learner('bricks-model.pkl')
+    categories = learn.dls.vocab
+    pred, idx, probs = learn.predict(img)
+    return {category: float(prob) for category, prob in zip(categories, probs)}
+
+def classify_image(img):
+    color_result = classify_image_color(img)
+    shape_result = classify_image_shape(img)
+    result = {}
+    for key in set(color_result.keys()) | set(shape_result.keys()):
+        result[key] = {"color": color_result.get(key, 0.0), "shape": shape_result.get(key, 0.0)}
+    return result
+
 
 image = gr.inputs.Image(shape=(256, 256))
-label = gr.outputs.Label(num_top_classes=3)
+label = gr.outputs.Label()
+
+intf = gr.Interface(
+    fn=classify_image, 
+    inputs=image, 
+    outputs=label, 
+    examples="", 
+    title="Lego Brick Classifier", 
+    layout="vertical"
+)
+intf.launch()
+
 
-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)