Add mask guidance

README.md

@@ -1,6 +1,6 @@
 ---
 title: Sidewalks Seg
-emoji: 📊
+emoji: 🚶‍♀️
 colorFrom: indigo
 colorTo: indigo
 sdk: gradio

app.py

@@ -1,4 +1,5 @@
 import torch
+import numpy as np
 import gradio as gr
 import matplotlib.pyplot as plt
 from PIL import Image
@@ -6,47 +7,103 @@ from transformers import SamModel, SamProcessor
 
 device = 'cuda' if torch.cuda.is_available() else 'cpu'
 processor = SamProcessor.from_pretrained('facebook/sam-vit-base')
-model = SamModel.from_pretrained('hmdliu/sidewalks-seg-base')
+model = SamModel.from_pretrained('hmdliu/sidewalks-seg')
 model.to(device)
 
-def segment_image(image, threshold):
-    # init data
-    width, height = image.size
-    prompt = [0, 0, width, height]
+def segment_image(image, threshold, x_min, y_min, x_max, y_max):
+    # tolerate TIFF image input
+    image.save('image.png')
+    # init input data
+    prompt = [x_min, y_min, x_max, y_max]
     inputs = processor(image, input_boxes=[[prompt]], return_tensors='pt')
     # make prediction
     outputs = model(pixel_values=inputs['pixel_values'].to(device),
                     input_boxes=inputs['input_boxes'].to(device),
                     multimask_output=False)
     prob_map = torch.sigmoid(outputs.pred_masks.squeeze()).cpu().detach()
-    prediction = (prob_map > threshold).float()
-    prob_map, prediction = prob_map.numpy(), prediction.numpy()
+    pred_mask = (prob_map > threshold).float().numpy()
     # visualize results
-    save_image(image, 'image.png')
-    save_image(prob_map, 'prob.png', cmap='jet')
-    save_image(prediction, 'mask.png', cmap='gray')
-    return Image.open('image.png'), Image.open('mask.png'), Image.open('prob.png')
+    plt.figure(figsize=(8, 8))
+    plt.imshow(prob_map.numpy(), cmap='jet', interpolation='nearest')
+    plt.axis('off')
+    plt.tight_layout()
+    plt.savefig('prob.png', bbox_inches='tight', pad_inches=0)
+    plt.close()
+    # post-processing
+    ret_image = Image.open('image.png')
+    ret_pred = (Image.open('image.png'), [(pred_mask, 'Sidewalks')])
+    ret_prob = Image.open('prob.png')
+    return ret_image, ret_pred, ret_prob
 
-def save_image(image, path, **kwargs):
+def segment_image_with_guidance(image, threshold, offset, x_min, y_min, x_max, y_max):
+    # tolerate TIFF image input
+    image['background'].save('image.png')
+    # init input data
+    prompt = [x_min, y_min, x_max, y_max]
+    img = Image.open('image.png').convert('RGB')
+    inputs = processor(img, input_boxes=[[prompt]], return_tensors='pt')
+    # make prediction
+    outputs = model(pixel_values=inputs['pixel_values'].to(device),
+                    input_boxes=inputs['input_boxes'].to(device),
+                    multimask_output=False)
+    prob_map = torch.sigmoid(outputs.pred_masks.squeeze()).cpu().detach()
+    # perform mask guidance
+    guidance_mask = (np.max(np.array(image['layers'][0]), axis=2) != 0).astype(float)
+    enhance_map = prob_map.numpy() + offset * guidance_mask
+    pred_mask = (enhance_map > threshold).astype(float)
+    # visualize results
     plt.figure(figsize=(8, 8))
-    plt.imshow(image, interpolation='nearest', **kwargs)
+    plt.imshow(enhance_map, cmap='jet', interpolation='nearest')
     plt.axis('off')
     plt.tight_layout()
-    plt.savefig(path, bbox_inches='tight', pad_inches=0)
+    plt.savefig('prob.png', bbox_inches='tight', pad_inches=0)
     plt.close()
+    # post-processing
+    regions = [(guidance_mask, 'Guidance'), (pred_mask, 'Sidewalks')]
+    return (image['background'], regions), Image.open('prob.png')
 
 with gr.Blocks() as demo:
-    with gr.Row():
-        with gr.Column():
-            image_input = gr.Image(type='pil', label='TIFF Image')
-            threshold_slider = gr.Slider(minimum=0, maximum=1, step=0.01, value=0.5, label='Prediction Threshold')
-            segment_button = gr.Button('Segment')
-        with gr.Column():
-            prediction = gr.Image(type='pil', label='Segmentation Result')
-            prob_map = gr.Image(type='pil', label='Probability Map')
-    segment_button.click(
+    with gr.Tab('Baseline'):
+        with gr.Row():
+            with gr.Column():
+                t1_input = gr.Image(type='pil', label='Input Image')
+                with gr.Row():
+                    t1_x_min = gr.Textbox(value=0, label='x_min')
+                    t1_y_min = gr.Textbox(value=0, label='y_min')
+                    t1_x_max = gr.Textbox(value=256, label='x_max')
+                    t1_y_max = gr.Textbox(value=256, label='y_max')
+                t1_slider = gr.Slider(minimum=0, maximum=1, step=0.01, value=0.5, label='Prediction Threshold')
+                t1_segment = gr.Button('Segment')
+            with gr.Column():
+                t1_pred = gr.AnnotatedImage(color_map={'Sidewalks': '#0000FF'}, label='Prediction')
+            with gr.Column():
+                t1_prob_map = gr.Image(type='pil', label='Probability Map')
+    with gr.Tab('Mask Guidance'):
+        with gr.Row():
+            with gr.Column():
+                t2_input = gr.ImageEditor(type='pil', crop_size='2:1', label='Input Image',
+                                          brush=gr.Brush(default_size='5', color_mode='fixed'),
+                                          sources=['upload'], transforms=[])
+                with gr.Row():
+                    t2_x_min = gr.Textbox(value=0, label='x_min')
+                    t2_y_min = gr.Textbox(value=0, label='y_min')
+                    t2_x_max = gr.Textbox(value=256, label='x_max')
+                    t2_y_max = gr.Textbox(value=256, label='y_max')
+                t2_thresh = gr.Slider(minimum=0, maximum=1, step=0.01, value=0.5, label='Prediction Threshold')
+                t2_offset = gr.Slider(minimum=0, maximum=1, step=0.01, value=0.4, label='Guidance Offset')
+                t2_segment = gr.Button('Segment')
+            with gr.Column():
+                t2_pred = gr.AnnotatedImage(color_map={'Guidance': '#FF0000', 'Sidewalks': '#0000FF'}, label='Prediction')
+            with gr.Column():
+                t2_prob_map = gr.Image(type='pil', label='Probability Map')
+    t1_segment.click(
         segment_image, 
-        inputs=[image_input, threshold_slider], 
-        outputs=[image_input, prediction, prob_map]
+        inputs=[t1_input, t1_slider, t1_x_min, t1_y_min, t1_x_max, t1_y_max], 
+        outputs=[t1_input, t1_pred, t1_prob_map]
+    )
+    t2_segment.click(
+        segment_image_with_guidance, 
+        inputs=[t2_input, t2_thresh, t2_offset, t2_x_min, t2_y_min, t2_x_max, t2_y_max], 
+        outputs=[t2_pred, t2_prob_map]
     )
 demo.launch(debug=True, show_error=True)