sdgf

app.py

@@ -2,6 +2,7 @@ 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')
@@ -9,6 +10,7 @@ def classify_image_color(img):
     pred, idx, probs = learn.predict(img)
     return {category: 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')
@@ -16,6 +18,7 @@ def classify_image_shape(img):
     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)
@@ -24,6 +27,7 @@ def classify_image(img):
         result[key] = {"color": color_result.get(key, 0.0), "shape": 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
@@ -33,11 +37,12 @@ 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"
+    fn=classify_image,
+    inputs=image,
+    outputs=label,
+    examples="",
+    title="Lego Brick Classifier",
+    layout="vertical",
+    postprocess=postprocess
 )
 intf.launch()