Update layout

app.py

@@ -22,18 +22,18 @@ def segment_image(image, threshold):
     prediction = (prob_map > threshold).float()
     prob_map, prediction = prob_map.numpy(), prediction.numpy()
     # visualize results
+    save_image(image, 'image.png')
+    save_image(prob_map, 'prob.png', cmap='jet')
+    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(prediction, cmap='gray', interpolation='nearest')
-    plt.axis('off')
-    plt.tight_layout()
-    plt.savefig('mask.png', bbox_inches='tight', pad_inches=0)
-    plt.figure(figsize=(8, 8))
-    plt.imshow(prob_map, cmap='jet', interpolation='nearest')
+    plt.imshow(image, interpolation='nearest', **kwargs)
     plt.axis('off')
     plt.tight_layout()
-    plt.savefig('heatmap.png', bbox_inches='tight', pad_inches=0)
+    plt.savefig(path, bbox_inches='tight', pad_inches=0)
     plt.close()
-    return Image.open('mask.png'), Image.open('heatmap.png')
 
 with gr.Blocks() as demo:
     with gr.Row():
@@ -43,11 +43,10 @@ with gr.Blocks() as demo:
             segment_button = gr.Button('Segment')
         with gr.Column():
             prediction = gr.Image(type='pil', label='Segmentation Result')
-        with gr.Column():
             prob_map = gr.Image(type='pil', label='Probability Map')
     segment_button.click(
         segment_image, 
         inputs=[image_input, threshold_slider], 
-        outputs=[prediction, prob_map]
+        outputs=[image_input, prediction, prob_map]
     )
 demo.launch(debug=True, show_error=True)