Update app.py

app.py

@@ -347,7 +347,7 @@ def switch_rows(v):
         print(frames[0])
         return frames
 
-def draw_mask(f, v, d, evt: gr.EventData):
+def draw_mask(f, n, v, d, evt: gr.EventData):
     print(v)
     vv = json.loads(v)
     pts = np.array(vv, np.int32)
@@ -360,7 +360,6 @@ def draw_mask(f, v, d, evt: gr.EventData):
     else:
         gradient = cv2.imread(f[2]).astype(np.uint8)
 
-    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)
     indices = np.arange(0,256)   # List of all colors 
@@ -376,8 +375,7 @@ def draw_mask(f, v, d, evt: gr.EventData):
     mask = cv2.inRange(diff, -int(n/2), int(n/2))
     d["layers"][0][mask>0] = (0,0,0,255)
 
-    bg = cv2.convertScaleAbs(palette[bg]) # Converting image back to uint8
-    gd = cv2.convertScaleAbs(palette[gd])
+    d["background"] = cv2.convertScaleAbs(palette[bg]) # Converting image back to uint8
     d["layers"][0] = cv2.fillPoly(d["layers"][0], [pts], (0,0,0,255))
     
     return gr.ImageEditor(value=d)
@@ -433,6 +431,7 @@ 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)
             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\"));
@@ -489,7 +488,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, mouse, output_mask], outputs=[output_mask])
+            output_mask.apply(fn=draw_mask, inputs=[gradients, levels, 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')