Update app.py

app.py

@@ -56,10 +56,6 @@ 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 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):
     # rgb_path = os.path.join(test_image_dir,img)
@@ -74,12 +70,13 @@ def gen_pred(img=inputs, model=seg_model):
     pred = model.predict(img)
     pred = np.squeeze(pred, axis=0)
     print(pred)
-    pred = convert_from_cv2_to_image(pred)
+    cv2_img = cv2.cvtColor(pred, cv2.COLOR_BGR2RGB)
+    pil_img = Image.fromarray(cv2_img)
     # color_coverted = cv2.cvtColor(pred, cv2.COLOR_BGR2RGB)
     # pil_image = Image.fromarray(pred)
     # PIL_image = Image.fromarray(pred.astype('uint8'), 'RGB')
     # return "UI in developing process ..."
-    return pred
+    return pil_img
 
 title = "<h1 style='text-align: center;'>Semantic Segmentation</h1>"
 description = "Upload an image and get prediction mask"