Update app.py

app.py

@@ -6,6 +6,7 @@ import urllib.request
 import numpy as np
 import keras.backend as K
 from transformers import pipeline
+from matplotlib import pyplot as plt
 
 from PIL import Image
 from matplotlib import cm
@@ -65,8 +66,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)
+    fig = plt.figure(figsize=(10, 7))
+    fig.add_subplot(rows, columns, 1)
+    plt.imshow(pred, interpolation='catrom')
+    plt.axis('off')
+    plt.title("Prediction")
+    plt.show()
     # return "UI in developing process ..."
-    return pred
+    return fig
 
 title = "<h1 style='text-align: center;'>Semantic Segmentation</h1>"
 description = "Upload an image and get prediction mask"