add result display panels for binary mask and cropped image

app.py

@@ -9,6 +9,7 @@ from PIL import ImageDraw
 from utils.tools_gradio import fast_process
 import copy
 import argparse
+from PIL import Image
 
 # Use ONNX to speed up the inference.
 ENABLE_ONNX = False
@@ -108,7 +109,7 @@ def reset(session_state):
     session_state['ori_image'] = None
     session_state['image_with_prompt'] = None
     session_state['feature'] = None
-    return None, session_state
+    return None, None, None, session_state
 
 
 def reset_all(session_state):
@@ -118,7 +119,7 @@ def reset_all(session_state):
     session_state['ori_image'] = None
     session_state['image_with_prompt'] = None
     session_state['feature'] = None
-    return None, None, session_state
+    return None, None, None, None, None, None, session_state
 
 
 def clear(session_state):
@@ -126,7 +127,7 @@ def clear(session_state):
     session_state['label_list'] = []
     session_state['box_list'] = []
     session_state['image_with_prompt'] = copy.deepcopy(session_state['ori_image'])
-    return session_state['ori_image'], session_state
+    return session_state['ori_image'], None, None, session_state
 
 
 def on_image_upload(
@@ -150,7 +151,7 @@ def on_image_upload(
     nd_image = np.array(image)
     session_state['feature'] = predictor.set_image(nd_image)
 
-    return image, session_state
+    return image, None, None, session_state
 
 
 def convert_box(xyxy):
@@ -229,7 +230,11 @@ def segment_with_points(
         withContours=withContours,
     )
 
-    return seg, session_state
+    binary_mask = np.where(annotations[0] > 0.5, 255, 0).astype(np.uint8)
+    mask = Image.fromarray(binary_mask)
+    binary_mask = np.expand_dims(binary_mask, axis=2)
+    crop = Image.fromarray(np.concatenate((session_state['ori_image'], binary_mask), axis=2), "RGBA")
+    return seg, mask, crop, session_state
 
 
 def segment_with_box(
@@ -300,13 +305,23 @@ def segment_with_box(
             use_retina=use_retina,
             withContours=withContours,
         )
-        return seg, session_state
-    return image, session_state
+        binary_mask = np.where(annotations[0] > 0.5, 255, 0).astype(np.uint8)
+        mask = Image.fromarray(binary_mask)
+        binary_mask = np.expand_dims(binary_mask, axis=2)
+        crop = Image.fromarray(np.concatenate((session_state['ori_image'], binary_mask), axis=2), "RGBA")
+        return seg, mask, crop, session_state
+    return image, None, None, session_state
 
 
 img_p = gr.Image(label="Input with points", type="pil")
 img_b = gr.Image(label="Input with box", type="pil")
 
+mask_p = gr.Image(label="Mask", type="pil", interactive=False)
+crop_p = gr.Image(label="Cropped image", type="pil", interactive=False)
+
+mask_b = gr.Image(label="Mask", type="pil", interactive=False)
+crop_b = gr.Image(label="Cropped image", type="pil", interactive=False)
+
 with gr.Blocks(css=css, title="EdgeSAM") as demo:
     session_state = gr.State({
         'coord_list': [],
@@ -339,7 +354,8 @@ with gr.Blocks(css=css, title="EdgeSAM") as demo:
                         clear_btn_p = gr.Button("Clear", variant="secondary")
                         reset_btn_p = gr.Button("Reset", variant="secondary")
                 with gr.Row():
-                    gr.Markdown(description_p)
+                    mask_p.render()
+                    crop_p.render()
 
         with gr.Row():
             with gr.Column():
@@ -347,11 +363,13 @@ with gr.Blocks(css=css, title="EdgeSAM") as demo:
                 gr.Examples(
                     examples=examples,
                     inputs=[img_p, session_state],
-                    outputs=[img_p, session_state],
+                    outputs=[img_p, mask_p, crop_p, session_state],
                     examples_per_page=8,
                     fn=on_image_upload,
                     run_on_click=True
                 )
+            with gr.Column():
+                gr.Markdown(description_p)
 
     with gr.Tab("Box mode") as tab_b:
         # Images
@@ -362,7 +380,9 @@ with gr.Blocks(css=css, title="EdgeSAM") as demo:
                 with gr.Column():
                     clear_btn_b = gr.Button("Clear", variant="secondary")
                     reset_btn_b = gr.Button("Reset", variant="secondary")
-                    gr.Markdown(description_b)
+                    with gr.Row():
+                        mask_b.render()
+                        crop_b.render()
 
         with gr.Row():
             with gr.Column():
@@ -370,30 +390,32 @@ with gr.Blocks(css=css, title="EdgeSAM") as demo:
                 gr.Examples(
                     examples=examples,
                     inputs=[img_b, session_state],
-                    outputs=[img_b, session_state],
+                    outputs=[img_b, mask_b, crop_b, session_state],
                     examples_per_page=8,
                     fn=on_image_upload,
                     run_on_click=True
                 )
+            with gr.Column():
+                gr.Markdown(description_b)
 
     with gr.Row():
         with gr.Column(scale=1):
             gr.Markdown(
                 "<center><img src='https://visitor-badge.laobi.icu/badge?page_id=chongzhou/edgesam' alt='visitors'></center>")
 
-    img_p.upload(on_image_upload, [img_p, session_state], [img_p, session_state])
-    img_p.select(segment_with_points, [add_or_remove, session_state], [img_p, session_state])
+    img_p.upload(on_image_upload, [img_p, session_state], [img_p, mask_p, crop_p, session_state])
+    img_p.select(segment_with_points, [add_or_remove, session_state], [img_p, mask_p, crop_p, session_state])
 
-    clear_btn_p.click(clear, [session_state], [img_p, session_state])
-    reset_btn_p.click(reset, [session_state], [img_p, session_state])
-    tab_p.select(fn=reset_all, inputs=[session_state], outputs=[img_p, img_b, session_state])
+    clear_btn_p.click(clear, [session_state], [img_p, mask_p, crop_p, session_state])
+    reset_btn_p.click(reset, [session_state], [img_p, mask_p, crop_p, session_state])
+    tab_p.select(fn=reset_all, inputs=[session_state], outputs=[img_p, mask_p, crop_p, img_b, mask_b, crop_b, session_state])
 
-    img_b.upload(on_image_upload, [img_b, session_state], [img_b, session_state])
-    img_b.select(segment_with_box, [session_state], [img_b, session_state])
+    img_b.upload(on_image_upload, [img_b, session_state], [img_b, mask_b, crop_b, session_state])
+    img_b.select(segment_with_box, [session_state], [img_b, mask_b, crop_b, session_state])
 
-    clear_btn_b.click(clear, [session_state], [img_b, session_state])
-    reset_btn_b.click(reset, [session_state], [img_b, session_state])
-    tab_b.select(fn=reset_all, inputs=[session_state], outputs=[img_p, img_b, session_state])
+    clear_btn_b.click(clear, [session_state], [img_b, mask_b, crop_b, session_state])
+    reset_btn_b.click(reset, [session_state], [img_b, mask_b, crop_b, session_state])
+    tab_b.select(fn=reset_all, inputs=[session_state], outputs=[img_p, mask_p, crop_p, img_b, mask_b, crop_b, session_state])
 
 demo.queue()
 # demo.launch(server_name=args.server_name, server_port=args.port)