Update app.py

app.py

@@ -70,13 +70,14 @@ def gen_pred(img=inputs, model=seg_model):
     pred = model.predict(img)
     pred = np.squeeze(pred, axis=0)
     print(pred)
-    cv2_img = cv2.cvtColor(pred, cv2.COLOR_BGR2RGB)
-    pil_img = Image.fromarray(cv2_img)
+    img_bytes = pred.tobytes()
+    nparr = np.frombuffer(img_bytes, np.byte)
+    pred_pil = cv2.imdecode(nparr, cv2.IMREAD_REDUCED_COLOR_8)
     # 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 pil_img
+    return pred_pil
 
 title = "<h1 style='text-align: center;'>Semantic Segmentation</h1>"
 description = "Upload an image and get prediction mask"