Update app.py

app.py

@@ -347,7 +347,7 @@ def switch_rows(v):
         print(frames[0])
         return frames
 
-def draw_mask(f, n, v, d, evt: gr.EventData):
+def draw_mask(f, n, c, v, d, evt: gr.EventData):
     print(v)
     vv = json.loads(v)
     pts = np.array(vv, np.int32)
@@ -372,10 +372,25 @@ def draw_mask(f, n, v, d, evt: gr.EventData):
     gd = cv2.cvtColor(gradient, cv2.COLOR_RGB2GRAY)
 
     diff = (bg-gd).astype(np.uint8)
-    mask = cv2.inRange(diff, -int(n/2), int(n/2))
+    mask = cv2.inRange(diff, -int(c/2), int(c/2))
     d["layers"][0][mask>0] = (0,0,0,255)
 
-    d["background"] = cv2.convertScaleAbs(palette[bg]) # Converting image back to uint8
+    bg = cv2.convertScaleAbs(palette[bg]) # Converting image back to uint
+
+    scale = 1
+    delta = 0
+    ddepth = cv2.CV_16S
+
+    grad = cv2.Sobel(bg, ddepth, 1, 1, ksize=1, scale=scale, delta=delta, borderType=cv2.BORDER_DEFAULT)
+    # Gradient-Y
+    # grad_y = cv.Scharr(gray,ddepth,0,1)
+    #grad_y = cv.Sobel(gray, ddepth, 0, 1, ksize=3, scale=scale, delta=delta, borderType=cv.BORDER_DEFAULT)
+    
+    abs_grad = cv2.convertScaleAbs(grad)
+    #abs_grad_y = cv.convertScaleAbs(grad_y)
+    #grad = cv.addWeighted(abs_grad_x, 0.5, abs_grad_y, 0.5, 0)
+
+    th, d["layers"][0] = cv2.threshold(abs_grad, 127, 255, cv2.THRESH_BINARY)
     d["layers"][0] = cv2.fillPoly(d["layers"][0], [pts], (0,0,0,255))
     
     return gr.ImageEditor(value=d)
@@ -431,7 +446,8 @@ with gr.Blocks(css=css) as demo:
             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")
             gradients = gr.Files(visible=False, value=['./gradient_large.png', './gradient.png', './gradient_small.png'])
-            levels = gr.Slider(label="Color levels", value=16, maximum=256, minimum=2, step=1)
+            levels = gr.Number(label="Color levels", value=16, maximum=32, minimum=2, step=1)
+            contour = gr.Number(label="Contour width", value=2, maximum=16, minimum=0, step=2)
             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\"));
@@ -488,7 +504,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=[gradients, levels, mouse, output_mask], outputs=[output_mask])
+            output_mask.apply(fn=draw_mask, inputs=[gradients, levels, contour, 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')