blocks example of segmentation with interactive sliders

app.py

@@ -64,7 +64,42 @@ feature_extractor, segmentation_model, segmentation_cfg = load_segmentation_mode
 pipe = load_diffusion_pipeline()
 pipe = pipe.to(device)
 
-# TODO: potentially use `gr.Gallery` to display different masks
+def fn_segmentation(image, max_kernel, min_kernel):
+    inputs = feature_extractor(images=image, return_tensors="pt")
+    outputs = segmentation_model(**inputs)
+
+    processed_sizes = torch.as_tensor(inputs["pixel_values"].shape[-2:]).unsqueeze(0)
+    result = feature_extractor.post_process_panoptic(outputs, processed_sizes)[0]
+
+    panoptic_seg = Image.open(io.BytesIO(result["png_string"])).resize((image.width, image.height))
+    panoptic_seg = np.array(panoptic_seg, dtype=np.uint8)
+
+    panoptic_seg_id = rgb_to_id(panoptic_seg)
+
+    raw_masks = []
+    for s in result['segments_info']:
+        m = panoptic_seg_id == s['id']
+        raw_masks.append(m.astype(np.uint8) * 255)
+    
+    masks = fn_clean(raw_masks, max_kernel, min_kernel)
+
+    return masks, raw_masks
+
+def fn_clean(masks, max_kernel, min_kernel):
+    out = []
+    for m in masks:
+        m = torch.FloatTensor(m)[None, None]
+        m = min_pool(m, min_kernel)
+        m = max_pool(m, max_kernel)
+        m = m.squeeze().numpy().astype(np.uint8)
+        out.append(m)
+
+    return out
+
+def fn_mask(image, mask_enabled):
+    if len(mask_enabled) == 0:
+        return image
+
 def fn_segmentation_diffusion(prompt, mask_indices, image, max_kernel, min_kernel, num_diffusion_steps):
     mask_indices = [int(i) for i in mask_indices.split(',')]
     inputs = feature_extractor(images=image, return_tensors="pt")
@@ -144,17 +179,46 @@ def fn_segmentation_diffusion(prompt, mask_indices, image, max_kernel, min_kerne
 
 # iface = gr.Series(
     # iface_segmentation, iface_diffusion,
-iface = gr.Interface(
-    fn=fn_segmentation_diffusion,
-    inputs=[
-        "text",
-        "text", 
-        gr.Image(value="http://images.cocodataset.org/val2017/000000039769.jpg", type='pil'),
-        gr.Slider(minimum=1, maximum=99, value=23, step=2),
-        gr.Slider(minimum=1, maximum=99, value=5, step=2),
-        gr.Slider(minimum=0, maximum=100, value=50, step=1),
-    ], 
-    outputs=[gr.Image(), gr.Image(), gr.Textbox(interactive=False)]
-)
-
-iface.launch()
+
+# iface = gr.Interface(
+    # fn=fn_segmentation_diffusion,
+    # inputs=[
+        # "text",
+        # "text", 
+        # gr.Image(value="http://images.cocodataset.org/val2017/000000039769.jpg", type='pil'),
+        # gr.Slider(minimum=1, maximum=99, value=23, step=2),
+        # gr.Slider(minimum=1, maximum=99, value=5, step=2),
+        # gr.Slider(minimum=0, maximum=100, value=50, step=1),
+    # ], 
+    # outputs=[gr.Image(), gr.Image(), gr.Textbox(interactive=False)]
+# )
+
+# iface = gr.Interface(
+    # fn=fn_segmentation,
+    # inputs=[
+        # gr.Image(value="http://images.cocodataset.org/val2017/000000039769.jpg", type='pil'),
+        # gr.Slider(minimum=1, maximum=99, value=23, step=2),
+        # gr.Slider(minimum=1, maximum=99, value=5, step=2),
+    # ],
+    # outputs=gr.Gallery()
+# )
+
+# iface.launch()
+
+demo = gr.Blocks()
+
+with demo:
+    input_image = gr.Image(value="http://images.cocodataset.org/val2017/000000039769.jpg", type='pil')
+    mask_gallery = gr.Gallery()
+    mask_storage = gr.State()
+
+    max_slider = gr.Slider(minimum=1, maximum=99, value=23, step=2)
+    min_slider = gr.Slider(minimum=1, maximum=99, value=5, step=2)
+
+    bt_masks = gr.Button("Compute Masks")
+
+    bt_masks.click(fn_segmentation, inputs=[input_image, max_slider, min_slider], outputs=[mask_gallery, mask_storage])
+    max_slider.change(fn_clean, inputs=[mask_storage, max_slider, min_slider], outputs=mask_gallery)
+    min_slider.change(fn_clean, inputs=[mask_storage, max_slider, min_slider], outputs=mask_gallery)
+
+demo.launch()