resizing mask to original image shape

app.py

@@ -13,7 +13,7 @@ model = tf.keras.models.load_model("./tf_model.h5")
 
 def predict(image):
     img = np.array(image)
-    shape = img.shape[:2]
+    original_shape = img.shape[:2]
 
     im = tf.image.resize(img, (128, 128))
     im = tf.cast(im, tf.float32) / 255.0
@@ -54,6 +54,7 @@ def predict(image):
         mask *= 255
 
     mask = np.array(Image.fromarray(mask).convert("L"))
+    mask = np.resize(mask, original_shape)
 
     return mask