fixed it, box_tensor was needed to be resized and the plot function was wonky

app.py

@@ -8,6 +8,7 @@ import torch
 from segment_anything import sam_model_registry
 import matplotlib.pyplot as plt
 from PIL import Image
+import torch.nn.functional as F
 import io
 
 def load_image(file_path):
@@ -24,11 +25,13 @@ def load_image(file_path):
     H, W = img.shape[:2]
     return img, H, W
 
+@torch.no_grad()
 def medsam_inference(medsam_model, img_embed, box_1024, H, W):
     box_torch = torch.as_tensor(box_1024, dtype=torch.float, device=img_embed.device)
     if len(box_torch.shape) == 2:
         box_torch = box_torch[:, None, :] # (B, 1, 4)
 
+    box_torch=box_torch.reshape(1,4) 
     sparse_embeddings, dense_embeddings = medsam_model.prompt_encoder(
         points=None,
         boxes=box_torch,
@@ -67,11 +70,7 @@ def visualize(image, mask, box):
     ax[1].imshow(image, cmap='gray')
     ax[1].imshow(mask, alpha=0.5, cmap="jet")
     plt.tight_layout()
-    buf = io.BytesIO()
-    plt.savefig(buf, format='png')
-    plt.close(fig)
-    buf.seek(0)
-    return buf
+    return fig 
 
 # Main function for Gradio app
 def process_images(file, x_min, y_min, x_max, y_max):
@@ -79,7 +78,11 @@ def process_images(file, x_min, y_min, x_max, y_max):
     
     # Load and preprocess image
     image, H, W = load_image(file)
-    image_resized = transform.resize(image, (1024, 1024), anti_aliasing=True)
+    if len(image.shape) == 2:
+        image = np.repeat(image[:, :, None], 3, axis=-1)
+    H, W, _ = image.shape
+
+    image_resized = transform.resize(image, (1024, 1024), order=3, preserve_range=True, anti_aliasing=True).astype(np.uint8)
     image_resized = (image_resized - image_resized.min()) / np.clip(image_resized.max() - image_resized.min(), a_min=1e-8, a_max=None)
     image_tensor = torch.tensor(image_resized).float().permute(2, 0, 1).unsqueeze(0).to(device)
 
@@ -102,7 +105,7 @@ def process_images(file, x_min, y_min, x_max, y_max):
 
     # Visualization
     visualization = visualize(image, mask, [x_min, y_min, x_max, y_max])
-    return visualization.getvalue()
+    return visualization #.getvalue()
 
 
 # Set up Gradio interface