Update app.py

app.py

@@ -70,9 +70,12 @@ def tif_bytes_to_pil_image(tif_bytes):
     # Create a BytesIO object from the TIFF bytes
     bytes_io = io.BytesIO(tif_bytes)
     
-    # Open the BytesIO object as an Image, resize to 256
+    # Open the BytesIO object as an Image, crop to square, resize to 256
     image = Image.open(bytes_io).convert("RGB")
-    #image = image.resize((256, 256))
+    w, h = image.size
+    dim = min(w, h)
+    image = image.crop(0, 0, dim, dim)
+    image = image.resize((256, 256))
 
     return image
 
@@ -94,7 +97,7 @@ def load_model():
 
     return model, processor, device
 
-def show_mask(mask, ax, width, height, random_color=False):
+def show_mask(mask, ax, random_color=False):
     if random_color:
         color = np.concatenate([np.random.random(3), np.array([0.6])], axis=0)
     else:
@@ -229,7 +232,7 @@ def server(input: Inputs, output: Outputs, session: Session):
         if image is None or prob is None or prediction is None:
             # Return a placeholder plot or message
             fig, ax = plt.subplots()
-            ax.text(0.5, 0.5, "Upload an image to see predictions", fontsize=12, ha="center")
+            ax.text(0.5, 0.5, "Upload a square image to see predictions. If the image is not a square, the image will be cropped to a square, taking the top left portion of the image. Predictions will take a few moments to load.", fontsize=12, ha="center")
             ax.axis("off")  # Hide axis
             plt.tight_layout()
             return fig
@@ -238,25 +241,22 @@ def server(input: Inputs, output: Outputs, session: Session):
         
         # Extract the image data
         #image_data = image.cpu().detach().numpy()
-
-        # Get the dimensions (width and height) of the image
-        width, height = image.size
         
         # 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, extent=(0, width, height, 0))  # Setting extent to match original image dimensions
-        axes[1].set_title(f"Probability Map {prob.shape}")
+        axes[1].imshow(prob)
+        axes[1].set_title("Probability Map")
         
         # Plot the prediction image on the right
-        axes[2].imshow(prediction, extent=(0, width, height, 0))  # Setting extent to match original image dimensions
+        axes[2].imshow(prediction) 
         axes[2].set_title("Prediction")
         
         # Plot the predicted mask on the right
         axes[3].imshow(image) 
-        show_mask(prediction, axes[3], width, height)
+        show_mask(prediction, axes[3])
         axes[3].set_title("Predicted Mask")
         
         # Hide axis ticks and labels