Update app.py

app.py

@@ -56,14 +56,9 @@ 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 convert_from_cv2_to_image(img: np.ndarray) -> Image:
+    # return Image.fromarray(img)
+    return Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
 
 
 def gen_pred(img=inputs, model=seg_model):
@@ -79,7 +74,7 @@ def gen_pred(img=inputs, model=seg_model):
     pred = model.predict(img)
     pred = np.squeeze(pred, axis=0)
     print(pred)
-    pred = Image.fromarray(np.uint8(cm.gist_earth(pred)*255))
+    pred = convert_from_cv2_to_image(pred)
     # color_coverted = cv2.cvtColor(pred, cv2.COLOR_BGR2RGB)
     # pil_image = Image.fromarray(pred)
     # PIL_image = Image.fromarray(pred.astype('uint8'), 'RGB')