<yxc

app.py

@@ -1,48 +1,41 @@
 import gradio as gr
+import requests
+import json
+from gradio.mix import Parallel
 from fastai.vision.all import load_learner
 
-
 def classify_image_color(img):
     learn = load_learner('model-color.pkl')
     categories = learn.dls.vocab
-    _, _, probs = learn.predict(img)
-    return [(category, float(prob)) for category, prob in zip(categories, probs)]
+    pred, idx, probs = learn.predict(img)
+    return {category: float(prob) for category, prob in zip(categories, probs)}
 
 def classify_image_shape(img):
     learn = load_learner('bricks-model.pkl')
     categories = learn.dls.vocab
-    _, _, probs = learn.predict(img)
-    return [(category, float(prob)) for category, prob in zip(categories, probs)]
+    pred, idx, probs = learn.predict(img)
+    return {category: float(prob) for category, prob in zip(categories, probs)}
 
 def classify_image(img):
-    color_result = dict(classify_image_color(img))
-    shape_result = dict(classify_image_shape(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
 
 def postprocess(prediction):
-    result = {}
-    for key, value in prediction.items():
-        result[key] = []
-        for inner_key, inner_value in value.items():
-            result[key].append((inner_key, inner_value))
-        result[key].sort(key=lambda x: x[1], reverse=True)
-    return result
-
-
+    return json.dumps(prediction)
 
 image = gr.inputs.Image(shape=(256, 256))
-label = gr.outputs.Label()
+output_json = gr.outputs.Textbox(type="auto", label="JSON Output")
 
 intf = gr.Interface(
     fn=classify_image, 
     inputs=image, 
-    outputs=label, 
+    outputs=output_json, 
     examples="", 
     title="Lego Brick Classifier", 
-    layout="vertical",
-    postprocess=postprocess
+    layout="vertical"
 )
-intf.launch()
+intf.launch(share=True)