Update app.py

app.py

@@ -251,7 +251,7 @@ def server(input: Inputs, output: Outputs, session: Session):
         axes[1].set_title(f"Probability Map {prob.shape}")
         
         # Plot the prediction image on the right
-        axes[2].imshow(prediction_im, extent=(0, width, height, 0))  # Setting extent to match original image dimensions
+        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