Update app.py

app.py

@@ -237,7 +237,7 @@ def server(input: Inputs, output: Outputs, session: Session):
             plt.tight_layout()
             return fig
         
-        fig, axes = plt.subplots(1, 4, constrained_layout=True)
+        fig, axes = plt.subplots(1, 4, , figsize=(15, 5))
         
         # Extract the image data
         #image_data = image.cpu().detach().numpy()
@@ -246,9 +246,9 @@ def server(input: Inputs, output: Outputs, session: Session):
         axes[0].imshow(image)  
         axes[0].set_title("Image")
         
-        # # Plot the probability map on the right
-        # axes[1].imshow(prob)
-        # axes[1].set_title("Probability Map")
+        # Plot the probability map on the right
+        axes[1].imshow(prob)
+        axes[1].set_title("Probability Map")
         
         # # Plot the prediction image on the right
         # axes[2].imshow(prediction)