Update app.py

app.py

@@ -16,12 +16,14 @@ def load_bounding_boxes(csv_file):
     df = pd.read_csv(csv_file)
     return df
 
-def load_dicom_image(filename):
-    if filename.endswith(".dcm"):
-        ds = pydicom.dcmread(filename)
+def load_image(file_path):
+    if file_path.endswith(".dcm"):
+        ds = pydicom.dcmread(file_path)
         img = ds.pixel_array
-        H, W = img.shape
-    return np.array(img), H, W
+    else:
+        img = np.array(Image.open(file_path).convert('L')) # Convert to grayscale
+    H, W = img.shape
+    return img, H, W
 
 # MedSAM inference function
 def medsam_inference(medsam_model, img, box, H, W, target_size):
@@ -59,12 +61,12 @@ def visualize(image, mask, box):
     return buf
 
 # Main function for Gradio app
-def process_images(dicom_file, x_min, y_min, x_max, y_max):
-    image, H, W = load_dicom_image(dicom_file)
+def process_images(file, x_min, y_min, x_max, y_max):
+    image, H, W = load_image(file)
 
     # Initialize MedSAM model
     device = 'cuda' if torch.cuda.is_available() else 'cpu'
-    medsam_model = sam_model_registry['vit_b'](checkpoint=MedSAM_CKPT_PATH) # Ensure the correct path
+    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()
 
@@ -78,7 +80,7 @@ def process_images(dicom_file, x_min, y_min, x_max, y_max):
 iface = gr.Interface(
     fn=process_images,
     inputs=[
-        gr.File(label="DICOM File"),
+        gr.File(label="MRI Slice (DICOM, PNG, etc.)"),
         gr.Number(label="X min"),
         gr.Number(label="Y min"),
         gr.Number(label="X max"),