Update app.py

app.py

@@ -336,16 +336,17 @@ def switch_rows(v):
         print(frames[0])
         return frames
 
-def draw_mask(v, d, evt: gr.EventData):
+def draw_mask(th, v, d, evt: gr.EventData):
     print(v)
     vv = json.loads(v)
     pts = np.array(vv, np.int32)
     pts = pts.reshape((-1,1,2))
 
     bg = cv2.cvtColor(d["background"], cv2.COLOR_RGB2GRAY)
-    gd = cv2.cvtColor(cv2.imread(masks[len(masks)-1]), cv2.COLOR_RGB2GRAY)-127
-
-    d["background"] = bg - gd
+    gd = (cv2.cvtColor(cv2.imread(masks[len(masks)-1]), cv2.COLOR_RGB2GRAY)-th).astype(np.uint8)
+    diff = abs(bg-gd).astype(np.uint8)
+    mask = cv2.inRange(diff, 0, 5)
+    d["background"][mask>0] = (0,0,0,255)
 
     d["layers"][0] = cv2.fillPoly(d["layers"][0], [pts], (0,0,0,255))
     d["composite"] = cv2.fillPoly(d["composite"], [pts], (0,0,0,255))
@@ -404,6 +405,7 @@ with gr.Blocks(css=css) as demo:
             output_switch = gr.Checkbox(label="Show depths")
             output_switch.input(fn=switch_rows, inputs=[output_switch], outputs=[output_frame])
             output_mask = gr.ImageEditor(layers=True, sources=None, type="numpy", interactive=True, transforms=(None,), eraser=gr.Eraser(), brush=gr.Brush(default_size=0, colors=['black', '#505050', '#a0a0a0', 'white']), elem_id="image_edit")
+            threshold = gr.Slider(value=127, minimum=1, maximum=255, step=1)
             selector = gr.HTML(value="""<div id='select_cursor' style='width:1px;height:1px;position:absolute;border:8px dotted #ffffff;opacity:50%;'></div>
             <a href='#' id='selector' onclick='this.style.fontWeight=\"bold\";
             document.body.appendChild(document.getElementById(\"select_cursor\"));
@@ -460,7 +462,7 @@ with gr.Blocks(css=css) as demo:
               document.getElementById(\"select_cursor\").style.top = 0;
             ' title='Clear selection' style='text-decoration:none;color:white;'>[x] Clear selection</a>""")
             mouse = gr.Textbox(elem_id="mouse", value="""[]""", interactive=False)
-            output_mask.apply(fn=draw_mask, inputs=[mouse, output_mask], outputs=[output_mask])
+            output_mask.apply(fn=draw_mask, inputs=[threshold, mouse, output_mask], outputs=[output_mask])
             submit = gr.Button("Submit")
         with gr.Column():
             model_type = gr.Dropdown([("small", "vits"), ("base", "vitb"), ("large", "vitl")], type="value", value="vits", label='Model Type')