Update app.py

app.py

@@ -347,14 +347,21 @@ def switch_rows(v):
         print(frames[0])
         return frames
 
-def draw_mask(th, v, d, evt: gr.EventData):
+def draw_mask(f, v, d, evt: gr.EventData):
     print(v)
     vv = json.loads(v)
     pts = np.array(vv, np.int32)
     pts = pts.reshape((-1,1,2))
 
+    if d["background"].shape[0] == 2048: #height
+        gradient = cv2.imread(f[0]).astype(np.uint8)
+    elif d["background"].shape[0] == 1024:
+        gradient = cv2.imread(f[1]).astype(np.uint8)
+    else:
+        gradient = cv2.imread(f[2]).astype(np.uint8)
+
     bg = cv2.cvtColor(d["background"], cv2.COLOR_RGB2GRAY)
-    gd = ((cv2.cvtColor(cv2.imread(masks[len(masks)-1]), cv2.COLOR_RGB2GRAY)-th)*2).astype(np.uint8)
+    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)
@@ -427,7 +434,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)
+            gradients = gr.Files(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\"));
@@ -484,7 +491,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=[threshold, mouse, output_mask], outputs=[output_mask])
+            output_mask.apply(fn=draw_mask, inputs=[gradients, 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')