Update app.py

app.py

@@ -91,6 +91,23 @@ def predict_on_train(image, mask):
 
     # Steps to get prediction of the satellite image
     mask_resized = resize_image(mask)
+
+    # Create a figure
+    fig, ax = plt.subplots()
+
+    # Display the image
+    ax.imshow(mask_resized)
+
+    # Convert the figure to a PIL Image
+    image_buffer = io.BytesIO()
+    plt.savefig(image_buffer, format='png')
+    image_buffer.seek(0)
+    mask_pil = Image.open(image_buffer)
+
+    # Close the figure to release resources
+    plt.close(fig)
+
+    # ----------------------------------------------
     
     # Steps to get prediction of the satellite image
     sample_image_resized = resize_image(image)
@@ -124,7 +141,7 @@ def predict_on_train(image, mask):
     # Close the figure to release resources
     plt.close(fig)
 
-    return mask_resized, image_pil
+    return mask_pil, image_pil
 
 def predict_on_test(image):