Swapped csv bb input for dialogue box

app.py

@@ -45,13 +45,12 @@ def medsam_inference(medsam_model, img, box, H, W, target_size):
     return mask_resized
 
 # Function for visualizing images with masks
-def visualize(images, masks, box):
-    fig, ax = plt.subplots(len(images), 2, figsize=(10, 5*len(images)))
-    for i, (image, mask) in enumerate(zip(images, masks)):
-        ax[i, 0].imshow(image, cmap='gray')
-        ax[i, 0].add_patch(plt.Rectangle((box[0], box[1]), box[2]-box[0], box[3]-box[1], edgecolor="red", facecolor="none"))
-        ax[i, 1].imshow(image, cmap='gray')
-        ax[i, 1].imshow(mask, alpha=0.5, cmap="jet")
+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()
     buf = io.BytesIO()
     plt.savefig(buf, format='png')
@@ -60,34 +59,32 @@ def visualize(images, masks, box):
     return buf
 
 # Main function for Gradio app
-def process_images(csv_file, dicom_file):
-    bounding_boxes = load_bounding_boxes(csv_file)
+def process_images(dicom_file, x_min, y_min, x_max, y_max):
     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['vit_b'](checkpoint="medsam_vit_b.pth") # Ensure the correct path
+    medsam_model = sam_model_registry['vit_b'](checkpoint=MedSAM_CKPT_PATH) # Ensure the correct path
     medsam_model = medsam_model.to(device)
     medsam_model.eval()
 
-    masks = []
-    boxes = []
-    for index, row in bounding_boxes.iterrows():
-        box = [row['x_min'], row['y_min'], row['x_max'], row['y_max']]
-        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)
+    box = [x_min, y_min, x_max, y_max]
+    mask = medsam_inference(medsam_model, image, box, H, W, H) # Assuming target size is the same as the image height
 
-    visualizations = visualize([image] * len(masks), masks, boxes)
-    return visualizations.getvalue()
+    visualization = visualize(image, mask, box)
+    return visualization.getvalue() # Returning the byte stream
 
 # Set up Gradio interface
 iface = gr.Interface(
     fn=process_images,
     inputs=[
-        gr.File(label="CSV File"), 
-        gr.File(label="DICOM File")],
-    outputs="plot"
+        gr.inputs.File(label="DICOM File"),
+        gr.inputs.Number(label="X min"),
+        gr.inputs.Number(label="Y min"),
+        gr.inputs.Number(label="X max"),
+        gr.inputs.Number(label="Y max")
+    ],
+    outputs=gr.outputs.Image(type="plot")
 )
 
 iface.launch()