changed interface to allow for bb selection

app.py

@@ -10,6 +10,7 @@ import matplotlib.pyplot as plt
 from PIL import Image
 import torch.nn.functional as F
 import io
+from gradio_image_prompter import ImagePrompter
 
 def load_image(file_path):
     if file_path.endswith(".dcm"):
@@ -38,10 +39,6 @@ def medsam_inference(medsam_model, img_embed, box_1024, H, W):
         masks=None,
     )
     
-    #print shapes of tensors 
-    print("Shape of sparse_embeddings:", sparse_embeddings.shape)
-    print("Shape of dense_embeddings:", dense_embeddings.shape)
-    
     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)
@@ -70,14 +67,28 @@ def visualize(image, mask, box):
     ax[1].imshow(image, cmap='gray')
     ax[1].imshow(mask, alpha=0.5, cmap="jet")
     plt.tight_layout()
-    return fig 
+
+    # 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(file, x_min, y_min, x_max, y_max):
+def process_images(img_dict):
     device = 'cuda' if torch.cuda.is_available() else 'cpu'
     
     # Load and preprocess image
-    image, H, W = load_image(file)
+    img = img_dict['image']
+    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]
+    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
@@ -105,11 +116,44 @@ def process_images(file, x_min, y_min, x_max, y_max):
 
     # Visualization
     visualization = visualize(image, mask, [x_min, y_min, x_max, y_max])
-    return visualization #.getvalue()
+    return visualization
+
+def echo(x_min, y_min, x_max, y_max):
+    print(x_min, y_min, x_max, y_max)
 
 
 # Set up Gradio interface
 iface = gr.Interface(
+    fn=process_images,
+    inputs=[
+        ImagePrompter(label="Select ROIs")  # Custom image prompter for selecting regions of interest
+    ],
+    outputs=[
+        gr.Image(type="pil", label="Processed Image"),  # Image output
+    ],
+    title="Image Processing with Custom Prompts",
+    description="Upload an image and select regions of interest for processing."
+)
+
+# Launch the interface
+iface.launch()
+
+'''iface= gr.Interface(fn=process_images,
+                    inputs=[lambda prompts: (prompts["image"], prompts["points"]),
+                    ImagePrompter(show_label=False)],
+                    outputs="plot")'''
+
+
+
+'''iface = gr.Interface(
+    lambda prompts: (prompts["image"], prompts["points"]),
+    ImagePrompter(show_label=False),
+    [gr.Image(show_label=False), gr.Dataframe(label="Points")],
+)
+'''
+
+
+'''gr.Interface(
     fn=process_images,
     inputs=[
         gr.File(label="MRI Slice (DICOM, PNG, etc.)"),
@@ -119,6 +163,4 @@ iface = gr.Interface(
         gr.Number(label="Y max")
     ],
     outputs="plot"
-)
-
-iface.launch()
+)'''