sdfgsdfg

app.py

@@ -2,32 +2,28 @@ import gradio as gr
 from gradio.mix import Parallel
 from fastai.vision.all import load_learner
 
-
 def classify_image_color(img):
     from fastai.vision.all import load_learner
     learn = load_learner('model-color.pkl')
     categories = learn.dls.vocab
     pred, idx, probs = learn.predict(img)
-    return {category: float(prob) for category, prob in zip(categories, probs)}
-
+    return {f"{category}_color": float(prob) for category, prob in zip(categories, probs)}
 
 def classify_image_shape(img):
     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)}
-
+    return {f"{category}_shape": 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)}
+        result[key] = color_result.get(key, 0.0) + shape_result.get(key, 0.0)
     return result
 
-
 def postprocess(prediction):
     sorted_pred = sorted(prediction.items(), key=lambda x: x[1], reverse=True)
     return sorted_pred
@@ -37,12 +33,11 @@ image = gr.inputs.Image(shape=(256, 256))
 label = gr.outputs.Label()
 
 intf = gr.Interface(
-    fn=classify_image,
-    inputs=image,
-    outputs=label,
-    examples="",
-    title="Lego Brick Classifier",
-    layout="vertical",
-    postprocess=postprocess
+    fn=classify_image, 
+    inputs=image, 
+    outputs=label, 
+    examples="", 
+    title="Lego Brick Classifier", 
+    layout="vertical"
 )
 intf.launch()