Update app.py

app.py

@@ -16,16 +16,12 @@ def load_bounding_boxes(csv_file):
     df = pd.read_csv(csv_file)
     return df
 
-# Function to load DICOM images
-def load_dicom_images(folder_path):
-    images = []
-    for filename in sorted(os.listdir(folder_path)):
-        if filename.endswith(".dcm"):
-            filepath = os.path.join(folder_path, filename)
-            ds = pydicom.dcmread(filepath)
-            img = ds.pixel_array
-            images.append(img)
-    return np.array(images)
+def load_dicom_image(filename):
+    if filename.endswith(".dcm"):
+        ds = pydicom.dcmread(filename)
+        img = ds.pixel_array
+        H, W = img.shape
+    return np.array(img), H, W
 
 # MedSAM inference function
 def medsam_inference(medsam_model, img, box, H, W, target_size):
@@ -64,42 +60,34 @@ def visualize(images, masks, box):
     return buf
 
 # Main function for Gradio app
-def process_images(csv_file, dicom_folder):
+def process_images(csv_file, dicom_file):
     bounding_boxes = load_bounding_boxes(csv_file)
-    dicom_images = load_dicom_images(dicom_folder)
+    image, H, W = load_dicom_image(dicom_file)
 
     # Initialize MedSAM model
     device = 'cuda' if torch.cuda.is_available() else 'cpu'
-    medsam_model = sam_model_registry['your_model_version'](checkpoint='path_to_your_checkpoint')
+    medsam_model = sam_model_registry['vit_b'](checkpoint="medsam_vit_b.pth") # Ensure the correct path
     medsam_model = medsam_model.to(device)
     medsam_model.eval()
 
     masks = []
+    boxes = []
     for index, row in bounding_boxes.iterrows():
-        if index >= len(dicom_images):
-            continue  # Skip if the index exceeds the number of images
-
-        image = dicom_images[index]
-        H, W = image.shape
         box = [row['x_min'], row['y_min'], row['x_max'], row['y_max']]
-
-        mask = medsam_inference(medsam_model, image, box, H, W, target_size)
+        mask = medsam_inference(medsam_model, image, box, H, W, H) # Assuming target size is the same as the image height
         masks.append(mask)
+        boxes.append(box)
 
-    visualizations = visualize(dicom_images, masks, box)
-
-    return visualizations, np.array(masks)
+    visualizations = visualize([image] * len(masks), masks, boxes)
+    return visualizations.getvalue()
 
 # Set up Gradio interface
 iface = gr.Interface(
     fn=process_images,
     inputs=[
         gr.File(label="CSV File"), 
-        gr.File(label="Zipped DICOM stack or nrrd file")],
-    outputs=[
-        gr.Image(type="pil"), 
-        gr.File(type="numpy")
-    ]
+        gr.File(label="DICOM File")],
+    outputs="plot"
 )
 
 iface.launch()