Update app.py

app.py

@@ -252,8 +252,9 @@ def server(input: Inputs, output: Outputs, session: Session):
         axes[1].set_title("Probability Map")
         
         # Plot the prediction image on the right
-        prediction = prediction[:height, :width]
-        axes[2].imshow(prediction)  # Setting extent to match original image dimensions
+        prediction = Image.fromarray(prediction)
+        prediction = prediction.crop((0,0,width,height))
+        axes[2].imshow(prediction, extent=(0, width, height, 0))  # Setting extent to match original image dimensions
         axes[2].set_title("Prediction")
         
         # Plot the predicted mask on the right