Update app.py

app.py

@@ -336,7 +336,7 @@ def switch_rows(v):
         print(frames[0])
         return frames
 
-def draw_mask(v, d, evt: gr.EventData):
+def draw_mask(ksize, th, v, d, evt: gr.EventData):
     print(v)
     vv = v.split(", ")
     x = int(vv[0])
@@ -349,16 +349,15 @@ def draw_mask(v, d, evt: gr.EventData):
     ddepth = cv2.CV_16S
 
     gray = cv2.cvtColor(d["background"], cv2.COLOR_BGR2GRAY)
-    grad_x = cv2.Sobel(gray, ddepth, 1, 0, ksize=5, scale=scale, delta=delta, borderType=cv2.BORDER_DEFAULT)
+    grad_x = cv2.Sobel(gray, ddepth, 1, 0, ksize=ksize, scale=scale, delta=delta, borderType=cv2.BORDER_DEFAULT)
     # Gradient-Y
     # grad_y = cv2.Scharr(gray,ddepth,0,1)
-    grad_y = cv2.Sobel(gray, ddepth, 0, 1, ksize=5, scale=scale, delta=delta, borderType=cv2.BORDER_DEFAULT)
+    grad_y = cv2.Sobel(gray, ddepth, 0, 1, ksize=ksize, scale=scale, delta=delta, borderType=cv2.BORDER_DEFAULT)
     abs_grad_x = cv2.convertScaleAbs(grad_x)
     abs_grad_y = cv2.convertScaleAbs(grad_y)
 
     grad = cv2.addWeighted(abs_grad_x, 0.5, abs_grad_y, 0.5, 0)
-    thresh = 127
-    bw = cv2.threshold(255-grad, thresh, 255, cv2.THRESH_BINARY)[1]
+    bw = cv2.threshold(255-grad, th, 255, cv2.THRESH_BINARY)[1]
 
     d["background"] = cv2.cvtColor(bw, cv2.COLOR_GRAY2RGB)
     d["composite"] = cv2.cvtColor(bw, cv2.COLOR_GRAY2RGB)
@@ -416,6 +415,8 @@ 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")
+            ksize = gr.Slider(label="Kernel", value=5, minimum=1, maximum=15, step=2)
+            threshold = gr.Slider(label="Threshold", value=127, minimum=0, maximum=255, step=1)
             selector = gr.HTML(value="""<div id='select_frame' style='position:absolute;border:1px dashed #808080;width:0;height:0'></div>
             <input type='number' id='mousedown' style='width:10em;color:auto;background-color:transparent;border:1px solid lightgray;'/>
             <input type='number' id='mouseup' style='width:10em;color:auto;background-color:transparent;border:1px solid lightgray;'/>
@@ -455,7 +456,7 @@ with gr.Blocks(css=css) as demo:
               document.getElementById(\"mouse\").getElementsByTagName(\"textarea\")[0].dispatchEvent(evt);
             }' title='select'>⬚</a>""")
             mouse = gr.Textbox(elem_id="mouse", value="""0, 0, 0, 0""")
-            output_mask.apply(fn=draw_mask, inputs=[mouse, output_mask], outputs=[output_mask])
+            output_mask.apply(fn=draw_mask, inputs=[ksize, 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')