Update app.py

app.py

@@ -247,10 +247,12 @@ def server(input: Inputs, output: Outputs, session: Session):
         axes[0].set_title("Image")
         
         # Plot the probability map on the right
+        prob = prob[:height, :width]
         axes[1].imshow(prob, extent=(0, width, height, 0))  # Setting extent to match original image dimensions
         axes[1].set_title("Probability Map")
         
         # Plot the prediction image on the right
+        prediction = prediction[:height, :width]
         axes[2].imshow(prediction, extent=(0, width, height, 0))  # Setting extent to match original image dimensions
         axes[2].set_title("Prediction")