Updated to allow for nrrd uploads

app.py

@@ -3,6 +3,7 @@ import pandas as pd
 import numpy as np
 import pydicom
 import os
+import nrrd
 from skimage import transform
 import torch
 from segment_anything import sam_model_registry
@@ -12,12 +13,18 @@ import torch.nn.functional as F
 import io
 from gradio_image_prompter import ImagePrompter
 
-def load_image(file_path):
-    if file_path.endswith(".dcm"):
-        ds = pydicom.dcmread(file_path)
-        img = ds.pixel_array
+def load_nrrd(file_path):
+    data, header = nrrd.read(file_path)
+    
+    # If the data is 3D, take the middle slice
+    if len(data.shape) == 3:
+        middle_slice = data.shape[2] // 2
+        img = data[:, :, middle_slice]
     else:
-        img = np.array(Image.open(file_path))
+        img = data
+
+    # Normalize the image to 0-255 range
+    img = ((img - img.min()) / (img.max() - img.min()) * 255).astype(np.uint8)
 
     # Convert grayscale to 3-channel RGB by replicating channels
     if len(img.shape) == 2:  # Grayscale image (height, width)
@@ -45,7 +52,7 @@ def medsam_inference(medsam_model, img_embed, box_1024, H, W):
         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)
 
@@ -59,7 +66,6 @@ def medsam_inference(medsam_model, img_embed, box_1024, H, W):
     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')
@@ -68,30 +74,24 @@ def visualize(image, mask, box):
     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
+    plt.close(fig)
     buf.seek(0)
     pil_img = Image.open(buf)
     
     return pil_img 
 
-# Main function for Gradio app
-def process_images(img_dict):
+def process_nrrd(nrrd_file, points):
     device = 'cuda' if torch.cuda.is_available() else 'cpu'
     
-    # Load and preprocess image
-    img = img_dict['image']
-    points = img_dict['points'][0]  # Accessing the first (and possibly only) set of points
+    # Load and preprocess NRRD file
+    image, H, W = load_nrrd(nrrd_file.name)
+    
     if len(points) >= 6:
         x_min, y_min, x_max, y_max = points[0], points[1], points[3], points[4]
     else:
         raise ValueError("Insufficient data for bounding box coordinates.")
-    image, H, W = img, img.shape[0], img.shape[1] #
-    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)
@@ -120,16 +120,17 @@ def process_images(img_dict):
 
 # Set up Gradio interface
 iface = gr.Interface(
-    fn=process_images,
+    fn=process_nrrd,
     inputs=[
-        ImagePrompter(label="Image")
+        gr.File(label="NRRD File"),
+        gr.JSON(label="Bounding Box Coordinates")
     ],
     outputs=[
         gr.Image(type="pil", label="Processed Image")
     ],
-    title="ROI Selection with MEDSAM",
-    description="Upload an image and select regions of interest for processing."
+    title="ROI Selection with MEDSAM for NRRD Files",
+    description="Upload an NRRD file and provide bounding box coordinates for processing."
 )
 
 # Launch the interface
-iface.launch()
+iface.launch()