Update app.py

app.py

@@ -247,13 +247,10 @@ 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")
+        axes[1].set_title(f"Probability Map {prob.shape}")
         
         # Plot the prediction image on the right
-        prediction_im = Image.fromarray(prediction)
-        prediction_im = prediction_im.crop((0,0,width,height))
         axes[2].imshow(prediction_im, extent=(0, width, height, 0))  # Setting extent to match original image dimensions
         axes[2].set_title("Prediction")