Update app.py

app.py

@@ -56,6 +56,16 @@ inputs = gr.inputs.Image(type="pil", label="Upload an image", source="upload")
 image_output = gr.outputs.Image(type="numpy", label="Output Image")
 # outputs = gr.outputs.HTML() #uncomment for single class output 
 
+def fig2img(fig):
+    """Convert a Matplotlib figure to a PIL Image and return it"""
+    import io
+    buf = io.BytesIO()
+    fig.savefig(buf)
+    buf.seek(0)
+    img = Image.open(buf)
+    return img
+
+
 def gen_pred(img=inputs, model=seg_model):
     # rgb_path = os.path.join(test_image_dir,img)
     # img = cv2.imread(rgb_path)
@@ -69,11 +79,12 @@ def gen_pred(img=inputs, model=seg_model):
     pred = model.predict(img)
     pred = np.squeeze(pred, axis=0)
     print(pred)
+    pred = fig2img(pred)
     # color_coverted = cv2.cvtColor(pred, cv2.COLOR_BGR2RGB)
     # pil_image = Image.fromarray(pred)
-    PIL_image = Image.fromarray(pred.astype('uint8'), 'RGB')
+    # PIL_image = Image.fromarray(pred.astype('uint8'), 'RGB')
     # return "UI in developing process ..."
-    return PIL_image
+    return pred
 
 title = "<h1 style='text-align: center;'>Semantic Segmentation</h1>"
 description = "Upload an image and get prediction mask"