upgrade to gradio4

README.md

@@ -4,19 +4,19 @@ emoji: 🩻
 colorFrom: blue
 colorTo: pink
 sdk: gradio
-sdk_version: 3.41.0
+sdk_version: 4.41.0
 app_file: app.py
 pinned: true
 license: apache-2.0
 ---
 
-This demo uses the ScribblePrompt-UNet model described in ["ScribblePrompt: Fast and Flexible Interactive Segmentation for Any Medical Image"](https://arxiv.org/abs/2312.07381)
+This demo uses the ScribblePrompt-UNet model described in ["ScribblePrompt: Fast and Flexible Interactive Segmentation for Any Biomedical Image"](https://arxiv.org/abs/2312.07381)
 
 ```
-@article{wong2023scribbleprompt,
-  title={ScribblePrompt: Fast and Flexible Interactive Segmentation for Any Medical Image},
+@article{wong2024scribbleprompt,
+  title={ScribblePrompt: Fast and Flexible Interactive Segmentation for Any Biomedical Image},
   author={Hallee E. Wong and Marianne Rakic and John Guttag and Adrian V. Dalca},
-  journal={arXiv:2312.07381},
-  year={2023},
+  journal={European Conference on Computer Vision (ECCV)},
+  year={2024},
 }
 ```

app.py

@@ -5,20 +5,19 @@ import torch.nn.functional as F
 import os
 import cv2
 import pathlib
+import math
 
 device = 'cuda' if torch.cuda.is_available() else 'cpu'
 
 from predictor import Predictor
 
-RES = 256
+H = 256
+W = 256
 
 test_example_dir = pathlib.Path("./test_examples")
 test_examples = [str(test_example_dir / x) for x in sorted(os.listdir(test_example_dir))]
 
-val_example_dir = pathlib.Path("./val_od_examples")
-val_examples = [str(val_example_dir / x) for x in sorted(os.listdir(val_example_dir))]
-
-default_example = test_example_dir / "TotalSegmentator_2.jpg" 
+default_example = test_examples[0]
 exp_dir = pathlib.Path('./checkpoints')
 default_model = 'ScribblePrompt-Unet'
 
@@ -82,7 +81,7 @@ def image_overlay(img, mask=None, scribbles=None, contour=False, alpha=0.5):
         
         if contour:
             contours = cv2.findContours((mask[...,None]>0.5).astype(np.uint8), cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE)
-            cv2.drawContours(output, contours[0], -1, (0, 255, 0), 1)
+            cv2.drawContours(output, contours[0], -1, (0, 255, 0), 2)
         else:
             mask_overlay = _get_overlay(img, mask)
             mask2 = 0.5*np.repeat(mask[...,None], 3, axis=-1)
@@ -111,26 +110,29 @@ def viz_pred_mask(img, mask=None, point_coords=None, point_labels=None, bbox_coo
 
     out = image_overlay(img, mask=mask, scribbles=seperate_scribble_masks)
 
+    H,W = img.shape[:2]
+    marker_size = min(H,W)//100
+
     if point_coords is not None:
         for i,(col,row) in enumerate(point_coords):
             if point_labels[i] == 1:
-                cv2.circle(out,(col, row), 2, (0,255,0), -1)
+                cv2.circle(out,(col, row), marker_size, (0,255,0), -1)
             else:
-                cv2.circle(out,(col, row), 2, (255,0,0), -1)
+                cv2.circle(out,(col, row), marker_size, (255,0,0), -1)
 
     if bbox_coords is not None:
         for i in range(len(bbox_coords)//2):
-            cv2.rectangle(out, bbox_coords[2*i], bbox_coords[2*i+1], (255,165,0), 1)
+            cv2.rectangle(out, bbox_coords[2*i], bbox_coords[2*i+1], (255,165,0), marker_size)
         if len(bbox_coords) % 2 == 1:
             cv2.circle(out, tuple(bbox_coords[-1]), 2, (255,165,0), -1)
 
-    return out
+    return out.astype(np.uint8)
 
 # -----------------------------------------------------------------------------
 # Collect scribbles
 # -----------------------------------------------------------------------------
 
-def get_scribbles(seperate_scribble_masks, last_scribble_mask, scribble_img, label: int):
+def get_scribbles(seperate_scribble_masks, last_scribble_mask, scribble_img):
     """
     Record scribbles
     """
@@ -138,28 +140,19 @@ def get_scribbles(seperate_scribble_masks, last_scribble_mask, scribble_img, lab
 
     if scribble_img is not None:
         
-        color_mask = scribble_img.get('mask')
-        scribble_mask = color_mask[...,0]/255
+        # Only use first layer
+        color_mask = scribble_img.get('layers')[0]
         
-        not_same = (scribble_mask != last_scribble_mask)
-        if not isinstance(not_same, bool):
-            not_same = not_same.any()
-
-        if not_same:
-            # In case any scribbles were removed
-            corrected_scribble_masks = np.stack(2*[(scribble_mask > 0)], axis=0)*seperate_scribble_masks
-            corrected_last_scribble_mask = last_scribble_mask*(scribble_mask > 0)
-
-            delta = (scribble_mask - corrected_last_scribble_mask) > 0
-            new_scribbles = scribble_mask * delta
-            corrected_scribble_masks[label,...] = np.clip(corrected_scribble_masks[label,...] + new_scribbles, a_min=0, a_max=1)
-
-            last_scribble_mask = scribble_mask
-            seperate_scribble_masks = corrected_scribble_masks
+        positive_scribbles = 1.0*(color_mask[...,1] > 128)
+        negative_scribbles = 1.0*(color_mask[...,0] > 128)
+        
+        seperate_scribble_masks = np.stack([positive_scribbles, negative_scribbles], axis=0)
+        last_scribble_mask = None
 
         return seperate_scribble_masks, last_scribble_mask
 
-def get_predictions(predictor, input_img, click_coords, click_labels, bbox_coords, seperate_scribble_masks, low_res_mask, img_features, multimask_mode):
+def get_predictions(predictor, input_img, click_coords, click_labels, bbox_coords, seperate_scribble_masks, 
+                    low_res_mask, img_features, multimask_mode):
     """
     Make predictions
     """
@@ -194,8 +187,7 @@ def refresh_predictions(predictor, input_img, output_img, click_coords, click_la
     
     # Record any new scribbles
     seperate_scribble_masks, last_scribble_mask = get_scribbles(
-        seperate_scribble_masks, last_scribble_mask, scribble_img, 
-        label=(0 if brush_label == "Positive (green)" else 1) # current color of the brush
+        seperate_scribble_masks, last_scribble_mask, scribble_img
     )
 
     # Make prediction
@@ -206,12 +198,33 @@ def refresh_predictions(predictor, input_img, output_img, click_coords, click_la
     # Update input visualizations
     mask_to_viz = best_mask.numpy()
     click_input_viz = viz_pred_mask(input_img, mask_to_viz, click_coords, click_labels, bbox_coords, seperate_scribble_masks, binary_checkbox)
-    scribble_input_viz = viz_pred_mask(input_img, mask_to_viz, click_coords, click_labels, bbox_coords, None, binary_checkbox)
     
+    empty_channel = np.zeros(input_img.shape[:2]).astype(np.uint8)
+    full_channel = 255*np.ones(input_img.shape[:2]).astype(np.uint8)
+    gray_mask = (255*mask_to_viz).astype(np.uint8)
+
+    bg = viz_pred_mask(input_img, mask_to_viz, click_coords, click_labels, bbox_coords, None, binary_checkbox)
+    old_scribbles = scribble_img.get('layers')[0]
+
+    scribble_mask = 255*(old_scribbles > 0).any(-1)
+    
+    scribble_input_viz = {
+        "background": np.stack([bg[...,i] for i in range(3)]+[full_channel], axis=-1),
+        ["layers"][0]: [np.stack([
+            (255*seperate_scribble_masks[1]).astype(np.uint8), 
+            (255*seperate_scribble_masks[0]).astype(np.uint8), 
+            empty_channel, 
+            scribble_mask
+        ], axis=-1)],
+        "composite": np.stack([click_input_viz[...,i] for i in range(3)]+[empty_channel], axis=-1),
+    }
+    
+    mask_img = 255*(mask_to_viz[...,None].repeat(axis=2, repeats=3)>0.5) if binary_checkbox else mask_to_viz[...,None].repeat(axis=2, repeats=3)
+
     out_viz = [
         viz_pred_mask(input_img, mask_to_viz, point_coords=None, point_labels=None, bbox_coords=None, seperate_scribble_masks=None, binary=binary_checkbox),
         input_img,
-        255*(mask_to_viz[...,None].repeat(axis=2, repeats=3)>0.5) if binary_checkbox else mask_to_viz[...,None].repeat(axis=2, repeats=3),
+        mask_img,
     ]
     
     return click_input_viz, scribble_input_viz, out_viz, best_mask, low_res_mask, img_features, seperate_scribble_masks, last_scribble_mask
@@ -298,8 +311,8 @@ def undo_click(predictor, input_img, brush_label, bbox_label, best_mask, low_res
 with gr.Blocks(theme=gr.themes.Default(text_size=gr.themes.sizes.text_lg)) as demo:
     
     # State variables
-    seperate_scribble_masks = gr.State(np.zeros((2,RES,RES), dtype=np.float32)) 
-    last_scribble_mask = gr.State(np.zeros((RES,RES), dtype=np.float32)) 
+    seperate_scribble_masks = gr.State(np.zeros((2, H, W), dtype=np.float32)) 
+    last_scribble_mask = gr.State(np.zeros((H, W), dtype=np.float32)) 
 
     click_coords = gr.State([])
     click_labels = gr.State([])
@@ -312,10 +325,11 @@ with gr.Blocks(theme=gr.themes.Default(text_size=gr.themes.sizes.text_lg)) as de
     low_res_mask = gr.State(None) 
 
     gr.HTML("""\
-    <h1 style="text-align: center; font-size: 28pt;">ScribblePrompt: Fast and Flexible Interactive Segmention for Any Medical Image</h1>
-    <p style="text-align: center; font-size: large;"><a href="https://scribbleprompt.csail.mit.edu">ScribblePrompt</a> is an interactive segmentation tool designed to help users segment <b>new</b> structures in medical images using scribbles, clicks <b>and</b> bounding boxes.
-    </p>
-                           
+    <h1 style="text-align: center; font-size: 28pt;">ScribblePrompt: Fast and Flexible Interactive Segmention for Any Biomedical Image</h1>
+    <p style="text-align: center; font-size: large;">
+            <b>ScribblePrompt</b> is an interactive segmentation tool designed to help users segment <b>new</b> structures in medical images using scribbles, clicks <b>and</b> bounding boxes.
+        [<a href="https://arxiv.org/abs/2312.07381">paper</a> | <a href="https://scribbleprompt.csail.mit.edu">website</a> | <a href="https://github.com/halleewong/ScribblePrompt">code</a>]    
+    </p>                           
     """)
 
     with gr.Accordion("Open for instructions!", open=False): 
@@ -351,34 +365,42 @@ with gr.Blocks(theme=gr.themes.Default(text_size=gr.themes.sizes.text_lg)) as de
                            value="Positive (green)", label="Scribble/Click Label")
             bbox_label = gr.Checkbox(value=False, label="Bounding Box (2 clicks)")
         with gr.Column(scale=1):
+            
             binary_checkbox = gr.Checkbox(value=True, label="Show binary masks", visible=False)
             autopredict_checkbox = gr.Checkbox(value=True, label="Auto-update prediction on clicks")
-            gr.Markdown("<span style='color:orange'>Troubleshooting:</span> If the image does not fully load in the Scribbles tab, click 'Clear Scribbles' or 'Clear All Inputs' to reload (it make take multiple tries). If you encounter an <span style='color:orange'>error</span> try clicking 'Clear All Inputs'.")
+            with gr.Accordion("Troubleshooting tips", open=False): 
+                gr.Markdown("<span style='color:orange'>If you encounter an <span style='color:orange'>error</span> try clicking 'Clear All Inputs'.")
             multimask_mode = gr.Checkbox(value=True, label="Multi-mask mode", visible=False)
 
     with gr.Row():
         display_height = 500
 
+        green_brush = gr.Brush(colors=["#00FF00"], color_mode="fixed", default_size=2)
+        red_brush = gr.Brush(colors=["#FF0000"], color_mode="fixed", default_size=2)
+
         with gr.Column(scale=1):
             with gr.Tab("Scribbles"):
-                scribble_img = gr.Image(
+                scribble_img = gr.ImageEditor(
                     label="Input",
-                    brush_radius=2, 
-                    interactive=True,
-                    brush_color="#00FF00", 
-                    tool="sketch", 
-                    height=display_height,
+                    image_mode="RGB",
+                    brush=green_brush,
                     type='numpy',
-                    value=default_example, 
+                    value=default_example,  
+                    transforms=(),
+                    sources=(),
+                    show_download_button=True,
+                    # height=display_height
                 )
-                clear_scribble_button = gr.ClearButton([scribble_img], value="Clear Scribbles", variant="stop")
 
             with gr.Tab("Clicks/Boxes") as click_tab:
                 click_img = gr.Image(
                     label="Input",
                     type='numpy',
                     value=default_example, 
-                    height=display_height
+                    show_download_button=True,
+                    sources=(),
+                    container=True,
+                    # height=display_height-50
                 )
                 with gr.Row():
                     undo_click_button = gr.Button("Undo Last Click")
@@ -388,21 +410,20 @@ with gr.Blocks(theme=gr.themes.Default(text_size=gr.themes.sizes.text_lg)) as de
                 input_img = gr.Image(
                     label="Input",
                     image_mode="L",
-                    visible=True, 
                     value=default_example, 
-                    height=display_height
+                    # height=display_height
                 ) 
                 gr.Markdown("To upload your own image: click the `x` in the top right corner to clear the current image, then drag & drop")
 
         with gr.Column(scale=1):
             with gr.Tab("Output"):
                 output_img = gr.Gallery(
-                    label='Outputs', 
+                    label='Output', 
                     columns=1, 
                     elem_id="gallery",
                     preview=True, 
                     object_fit="scale-down",
-                    height=display_height+50
+                    # height=display_height
                 )
     
     submit_button = gr.Button("Refresh Prediction", variant='primary')
@@ -424,28 +445,9 @@ with gr.Blocks(theme=gr.themes.Default(text_size=gr.themes.sizes.text_lg)) as de
     
     gr.Examples(examples=test_examples,
                 inputs=[input_img],
-                examples_per_page=10,
-                label='Unseen Examples from Test Datasets'
-                )
-    
-    gr.Examples(examples=val_examples,
-                inputs=[input_img],
-                examples_per_page=10,
-                label='Unseen Examples from Validation Datasets'
+                examples_per_page=12,
+                label='Examples from datasets unseen during training'
                 )
-    
-    # When clear scribble button is clicked
-    def clear_scribble_history(input_img):
-        if input_img is not None:
-            input_shape = input_img.shape[:2]
-        else:
-            input_shape = (RES, RES)
-        return input_img, input_img, np.zeros((2,)+input_shape, dtype=np.float32), np.zeros(input_shape, dtype=np.float32), None, None
-
-    clear_scribble_button.click(clear_scribble_history, 
-        inputs=[input_img], 
-        outputs=[click_img, scribble_img, seperate_scribble_masks, last_scribble_mask, best_mask, low_res_mask]
-    )
 
     # When clear clicks button is clicked
     def clear_click_history(input_img):
@@ -460,9 +462,25 @@ with gr.Blocks(theme=gr.themes.Default(text_size=gr.themes.sizes.text_lg)) as de
         if input_img is not None:
             input_shape = input_img.shape[:2]
         else:
-            input_shape = (RES, RES)
+            input_shape = (H, W)
         return input_img, input_img, [], [], [], [], np.zeros((2,)+input_shape, dtype=np.float32), np.zeros(input_shape, dtype=np.float32), None, None, None
 
+    # def clear_history_and_pad_input(input_img):
+    #     if input_img is not None:
+    #         h,w = input_img.shape[:2]
+    #         if h != w:
+    #             # Pad to square
+    #             pad = abs(h-w)
+    #             if h > w:
+    #                 padding = [(0,0), (math.ceil(pad/2),math.floor(pad/2))]
+    #             else:
+    #                 padding = [(math.ceil(pad/2),math.floor(pad/2)), (0,0)]
+
+    #             input_img = np.pad(input_img, padding, mode='constant', constant_values=0)
+
+    #     return clear_all_history(input_img)
+
+    
     input_img.change(clear_all_history,
                     inputs=[input_img], 
                     outputs=[click_img, scribble_img, 
@@ -527,7 +545,8 @@ with gr.Blocks(theme=gr.themes.Default(text_size=gr.themes.sizes.text_lg)) as de
     undo_click_button.click(fn=undo_click, 
                             inputs=[
                                 predictor,
-                                input_img, brush_label, bbox_label, best_mask, low_res_mask, click_coords, click_labels, bbox_coords,
+                                input_img, brush_label, bbox_label, best_mask, low_res_mask, click_coords, 
+                                click_labels, bbox_coords,
                                 seperate_scribble_masks, last_scribble_mask, scribble_img, img_features,
                                 output_img, binary_checkbox, multimask_mode, autopredict_checkbox
                             ], 
@@ -542,8 +561,7 @@ with gr.Blocks(theme=gr.themes.Default(text_size=gr.themes.sizes.text_lg)) as de
         Draw scribbles in the click canvas
         """
         seperate_scribble_masks, last_scribble_mask = get_scribbles(
-            seperate_scribble_masks, last_scribble_mask, scribble_img, 
-            label=(0 if brush_label == "Positive (green)" else 1) # previous color of the brush
+            seperate_scribble_masks, last_scribble_mask, scribble_img
         )
         click_input_viz = viz_pred_mask(
             input_img, best_mask, click_coords, click_labels, bbox_coords, seperate_scribble_masks, binary_checkbox
@@ -566,17 +584,11 @@ with gr.Blocks(theme=gr.themes.Default(text_size=gr.themes.sizes.text_lg)) as de
         Recorn new scribbles when changing brush color
         """
         if label == "Negative (red)":
-            brush_update = gr.Image.update(brush_color = "#FF0000") # red
+            brush_update = gr.update(brush=red_brush)
         elif label == "Positive (green)":
-            brush_update = gr.Image.update(brush_color = "#00FF00") # green
+            brush_update = gr.update(brush=green_brush)
         else:
             raise TypeError("Invalid brush color")
-        
-        # Record latest scribbles
-        seperate_scribble_masks, last_scribble_mask = get_scribbles(
-            seperate_scribble_masks, last_scribble_mask, scribble_img, 
-            label=(1 if label == "Positive (green)" else 0) # previous color of the brush
-        )
 
         return seperate_scribble_masks, last_scribble_mask, brush_update
     

predictor.py

@@ -3,7 +3,6 @@ import torch.nn.functional as F
 from typing import Dict, Tuple, Optional
 import network
 
-
 class Predictor:
     """
     Wrapper for ScribblePrompt Unet model
@@ -96,6 +95,7 @@ def rescale_inputs(inputs: Dict[str,any], res=128):
     Rescale the inputs 
     """ 
     h,w = inputs['img'].shape[-2:]
+
     if h != res or w != res:
         
         inputs.update(dict(