Update app.py

app.py

@@ -360,12 +360,6 @@ def draw_mask(f, v, d, evt: gr.EventData):
     else:
         gradient = cv2.imread(f[2]).astype(np.uint8)
 
-    bg = cv2.cvtColor(d["background"], cv2.COLOR_RGB2GRAY)
-    gd = cv2.cvtColor(gradient, cv2.COLOR_RGB2GRAY)
-    diff = abs(bg-gd).astype(np.uint8)
-    mask = cv2.inRange(diff, 0, 5)
-    d["background"][mask>0] = (0,0,0,255)
-
     n = 16    # Number of levels of quantization
     #for i in range(n):
     #    im[(im >= i*255/n) & (im < (i+1)*255/n)] = i*255/(n-1)
@@ -375,7 +369,15 @@ def draw_mask(f, v, d, evt: gr.EventData):
     color_levels = np.clip(np.int0(indices/divider),0,n-1) # color levels 0,1,2..
     palette = quantiz[color_levels]
 
-    d["background"] = cv2.convertScaleAbs(palette[d["background"]]) # Converting image back to uint8
+    bg = cv2.cvtColor(d["background"], cv2.COLOR_RGB2GRAY)
+    gd = cv2.cvtColor(gradient, cv2.COLOR_RGB2GRAY)
+
+    bg = cv2.convertScaleAbs(palette[bg]) # Converting image back to uint8
+    gd = cv2.convertScaleAbs(palette[gd])
+
+    diff = (bg-gd).astype(np.uint8)
+    mask = cv2.inRange(diff, -int(n/2), int(n/2))
+    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))
@@ -434,7 +436,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")
-            gradients = gr.Files(value=['./gradient_large.png', './gradient.png', './gradient_small.png'])
+            gradients = gr.Files(visible=False, value=['./gradient_large.png', './gradient.png', './gradient_small.png'])
             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\"));