Update app.py

app.py

@@ -383,9 +383,13 @@ def switch_rows(v):
         print(frames[0])
         return frames
 
-def apply_mask(d):
+def apply_mask(d, b):
     global frame_selected
     global masks
+
+    dilation = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (2 * b + 1, 2 * b + 1), (b, b))
+    d["layers"][0] = cv2.dilate(d["layers"][0], dilation)
+    d["background"][d["layers"][0]==(255,255,255,255)] = (255,255,255,255)
     
     masks[frame_selected] = d["background"]
     cv2.imwrite(f"f{frame_selected}_dmap.png", masks[frame_selected])
@@ -446,8 +450,8 @@ def draw_mask(l, t, v, d, evt: gr.EventData):
     # 255 << 8 tells to fill with the value 255)
     mask = mask[1:mask.shape[0]-1, 1:mask.shape[1]-1]
     
-    d["layers"][0][mask>0] = (0,0,0,255)
-    d["background"][mask>0] = (0,0,0,255)
+    d["layers"][0][mask>0] = (255,255,255,255)
+    d["background"][mask>0] = (255,255,255,255)
     
     return gr.ImageEditor(value=d)
 
@@ -560,12 +564,13 @@ with gr.Blocks(css=css, js=js) as demo:
             output_switch.input(fn=switch_rows, inputs=[output_switch], outputs=[output_frame])
             output_mask = gr.ImageEditor(layers=False, sources=('upload', 'clipboard'), show_download_button=True, type="numpy", interactive=True, transforms=(None,), eraser=gr.Eraser(), brush=gr.Brush(default_size=0, colors=['black', '#505050', '#a0a0a0', 'white']), elem_id="image_edit")
             depth_file = gr.File(label="Edited depth")
-            output_mask.apply(fn=apply_mask, inputs=[output_mask], outputs=[output_mask, depth_file])
             with gr.Accordion(label="Edge", open=False):
               levels = gr.Slider(label="Color levels", value=16, maximum=32, minimum=2, step=1)
               tolerance = gr.Slider(label="Tolerance", value=1, maximum=7, minimum=0, step=1)
+              bsize = gr.Slider(label="Border size", value=1, maximum=256, minimum=1, step=2)
               mouse = gr.Textbox(elem_id="mouse", value="""[]""", interactive=False)
               mouse.input(fn=draw_mask, show_progress="minimal", inputs=[levels, tolerance, mouse, output_mask], outputs=[output_mask])
+              output_mask.apply(fn=apply_mask, inputs=[output_mask, bsize], outputs=[output_mask, depth_file])
             selector = gr.HTML(value="""
             <a href='#' id='selector' onclick='this.style.fontWeight=\"bold\";
             document.getElementById(\"image_edit\").getElementsByTagName(\"canvas\")[0].oncontextmenu = function(e){e.preventDefault();}