Update app.py

app.py

@@ -94,14 +94,14 @@ def load_model():
 
     return model, processor, device
 
-def show_mask(mask, ax, random_color=False):
+def show_mask(mask, ax, width, height, random_color=False):
     if random_color:
         color = np.concatenate([np.random.random(3), np.array([0.6])], axis=0)
     else:
         color = np.array([30/255, 144/255, 255/255, 0.6])
     h, w = mask.shape[-2:]
     mask_image = mask.reshape(h, w, 1) * color.reshape(1, 1, -1)
-    ax.imshow(mask_image)
+    ax.imshow(mask_image, extent=(0, width, height, 0))  # Setting extent to match original image dimensions
 
 ### SERVER ###
 def server(input: Inputs, output: Outputs, session: Session):
@@ -238,22 +238,25 @@ def server(input: Inputs, output: Outputs, session: Session):
         
         # Extract the image data
         #image_data = image.cpu().detach().numpy()
+
+        # Get the dimensions of the original image
+        height, width = image.shape[:2]
         
         # Plot the first image on the left
         axes[0].imshow(image)  
         axes[0].set_title("Image")
         
         # Plot the probability map on the right
-        axes[1].imshow(prob)
+        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
-        axes[2].imshow(prediction)
+        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
         axes[3].imshow(image) 
-        show_mask(prediction, axes[3])
+        show_mask(prediction, axes[3], width, height)
         axes[3].set_title("Predicted Mask")
         
         # Hide axis ticks and labels