updated app

app.py

@@ -26,6 +26,14 @@ def segment_sidewalk(image):
     save_image(prediction, 'mask.png', cmap='gray')
     return Image.open('image.png'), Image.open('mask.png'), Image.open('prob.png')
 
+def save_image(image, path, **kwargs):
+    plt.figure(figsize=(8, 8))
+    plt.imshow(image, interpolation='nearest', **kwargs)
+    plt.axis('off')
+    plt.tight_layout()
+    plt.savefig(path, bbox_inches='tight', pad_inches=0)
+    plt.close()
+
 with gr.Blocks() as demo:
     with gr.Row():
         with gr.Column():