restored inference functionality

app.py

@@ -29,15 +29,63 @@ def load_image(file_path):
     H, W = img.shape[:2]
     return img, H, W
 
-# The rest of the code remains the same...
+@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,
+        masks=None,
+    )
+    
+    low_res_logits, _ = medsam_model.mask_decoder(
+        image_embeddings=img_embed, # (B, 256, 64, 64)
+        image_pe=medsam_model.prompt_encoder.get_dense_pe(), # (1, 256, 64, 64)
+        sparse_prompt_embeddings=sparse_embeddings, # (B, 2, 256)
+        dense_prompt_embeddings=dense_embeddings, # (B, 256, 64, 64)
+        multimask_output=False,
+        )
+
+    low_res_pred = torch.sigmoid(low_res_logits)  # (1, 1, 256, 256)
+
+    low_res_pred = F.interpolate(
+        low_res_pred,
+        size=(H, W),
+        mode="bilinear",
+        align_corners=False,
+    )  # (1, 1, gt.shape)
+    low_res_pred = low_res_pred.squeeze().cpu().numpy()  # (256, 256)
+    medsam_seg = (low_res_pred > 0.5).astype(np.uint8)
+    return medsam_seg
+
+# Function for visualizing images with masks
+def visualize(image, mask, box):
+    fig, ax = plt.subplots(1, 2, figsize=(10, 5))
+    ax[0].imshow(image, cmap='gray')
+    ax[0].add_patch(plt.Rectangle((box[0], box[1]), box[2] - box[0], box[3] - box[1], edgecolor="red", facecolor="none"))
+    ax[1].imshow(image, cmap='gray')
+    ax[1].imshow(mask, alpha=0.5, cmap="jet")
+    plt.tight_layout()
+
+    # Convert matplotlib figure to a PIL Image
+    buf = io.BytesIO()
+    fig.savefig(buf, format='png')
+    plt.close(fig)  # Close the figure to release memory
+    buf.seek(0)
+    pil_img = Image.open(buf)
+    
+    return pil_img 
 # Main function for Gradio app
 def process_images(img_dict):
     device = 'cuda' if torch.cuda.is_available() else 'cpu'
     
     # Load and preprocess image
+    print(img_dict)
     img = img_dict['image']
-    print(image.type())
     points = img_dict['points'][0]  # Accessing the first (and possibly only) set of points
     if len(points) >= 6:
         x_min, y_min, x_max, y_max = points[0], points[1], points[3], points[4]
@@ -48,12 +96,36 @@ def process_images(img_dict):
         image = np.repeat(image[:, :, None], 3, axis=-1)
     H, W, _ = image.shape
 
-    # The rest of the function remains the same...
+    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)
 
+    # Initialize the MedSAM model and set the device
+    model_checkpoint_path = "medsam_vit_b.pth"  # Replace with the correct path to your checkpoint
+    medsam_model = sam_model_registry['vit_b'](checkpoint=model_checkpoint_path)
+    medsam_model = medsam_model.to(device)
+    medsam_model.eval()
+
+    # Generate image embedding
+    with torch.no_grad():
+        img_embed = medsam_model.image_encoder(image_tensor)
+
+    # Calculate resized box coordinates
+    scale_factors = np.array([1024 / W, 1024 / H, 1024 / W, 1024 / H])
+    box_1024 = np.array([x_min, y_min, x_max, y_max]) * scale_factors
+
+    # Perform inference
+    mask = medsam_inference(medsam_model, img_embed, box_1024, H, W)
+
+    # Visualization
+    visualization = visualize(image, mask, [x_min, y_min, x_max, y_max])
+    return visualization
 # Set up Gradio interface
 iface = gr.Interface(
     fn=process_images,
-    inputs=[gr.File(label="image or nrrd")],
+    inputs=[
+        ImagePrompter(label="Image")
+    ],
     outputs=[
         gr.Image(type="pil", label="Processed Image")
     ],
@@ -62,4 +134,4 @@ iface = gr.Interface(
 )
 
 # Launch the interface
-iface.launch()
+iface.launch()