Update app.py

app.py

@@ -64,17 +64,7 @@ def gen_pred(img, model=seg_model):
     img = img/255
     img = tf.expand_dims(img, axis=0)
     pred = model.predict(img)
-    # pred = np.squeeze(pred, axis=0)
-    # image = Image.fromarray(np.uint8(cm.plasma(pred)*255))
-    # pil_img = Image.fromarray(pred, 'RGB')
-
-    # importing Image from PIL
-    # img_tmp = cv2.cvtColor(pred, cv2.COLOR_BGR2RGB)
-    # im_PIL = Image.fromarray(img_tmp)
-    # 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)
+    pred = np.squeeze(pred, axis=0)
     # return "UI in developing process ..."
     return pred
 
@@ -84,7 +74,7 @@ description = "Upload an image and get prediction mask"
 
 gr.Interface(fn=gen_pred, 
              inputs=[gr.Image(type='pil')], 
-             outputs=gr.Image(type='pil'), 
+             outputs=["plot"], 
              title=title, 
              examples=[["003e2c95d.jpg"], ["003b50a15.jpg"], ["003b48a9e.jpg"], ["0038cbe45.jpg"], ["00371aa92.jpg"]],
              description=description,