Update app.py

app.py

@@ -11,6 +11,7 @@ from PIL import Image
 from matplotlib import cm
 from tensorflow import keras
 
+from matplotlib import cm
 
 resized_shape = (768, 768, 3)
 IMG_SCALING = (1, 1)
@@ -64,14 +65,14 @@ def gen_pred(img, model=seg_model):
     img = tf.expand_dims(img, axis=0)
     pred = model.predict(img)
     pred = np.squeeze(pred, axis=0)
-    pil_img = Image.fromarray((pred * 255).astype(np.uint8))
+    image = Image.fromarray(np.uint8(cm.plasma(pred)*255))
     # pil_img = Image.fromarray(pred, 'RGB')
     # im = Image.fromarray((pred * 255).astype(np.uint8))
     # img_bytes = pred.tobytes()
     # nparr = np.frombuffer(img_bytes, np.byte)
     # pred_pil = cv2.imdecode(nparr, cv2.IMREAD_REDUCED_COLOR_8)
     # return "UI in developing process ..."
-    return pil_img
+    return image
 
 title = "<h1 style='text-align: center;'>Semantic Segmentation</h1>"
 description = "Upload an image and get prediction mask"