Update app.py

app.py

@@ -383,6 +383,15 @@ def switch_rows(v):
         print(frames[0])
         return frames
 
+def reset_mask():
+    global frame_selected
+    global masks
+    global depths
+    
+    masks[frame_selected] = depths[frame_selected]
+    cv2.imwrite(f"f{frame_selected}_dmap.png", masks[frame_selected])
+    return masks[frame_selected], f"f{frame_selected}_dmap.png"
+
 def apply_mask(d, b):
     global frame_selected
     global masks
@@ -561,11 +570,13 @@ with gr.Blocks(css=css, js=js) as demo:
             input_url = gr.Textbox(elem_id="url_in", value="./examples/streetview.mp4", label="URL")
             input_video = gr.Video(label="Input Video", format="mp4")
             input_url.input(fn=loadfile, inputs=[input_url], outputs=[input_video])
+            submit = gr.Button("Submit")
             output_frame = gr.Gallery(label="Frames", preview=True, columns=8192, type="numpy")
             output_switch = gr.Checkbox(label="Show depths")
             output_switch.input(fn=switch_rows, inputs=[output_switch], outputs=[output_frame])
             output_mask = gr.ImageEditor(layers=False, sources=('upload', 'clipboard'), show_download_button=True, type="numpy", interactive=True, transforms=(None,), eraser=gr.Eraser(), brush=gr.Brush(default_size=0, colors=['black', '#505050', '#a0a0a0', 'white']), elem_id="image_edit")
-            selector = gr.HTML(value="""
+            with gr.Row():
+              selector = gr.HTML(value="""
             <a href='#' id='selector' onclick='if (this.style.fontWeight!=\"bold\") {
             this.style.fontWeight=\"bold\";
             document.getElementById(\"image_edit\").getElementsByTagName(\"canvas\")[0].oncontextmenu = function(e){e.preventDefault();}
@@ -602,8 +613,9 @@ with gr.Blocks(css=css, js=js) as demo:
               
               document.getElementById(\"mouse\").getElementsByTagName(\"textarea\")[0].value = \"[]\";
             
-            ' title='Clear selection' style='text-decoration:none;color:white;'>✕ Clear selection</a>""")
-            apply = gr.Button("Apply")
+            ' title='Clear selection' style='text-decoration:none;color:white;'>✕ Clear</a>""")
+              apply = gr.Button("Apply", size='sm')
+              reset = gr.Button("Reset", size='sm')
             depth_file = gr.File(label="Edited depth")
             with gr.Accordion(label="Edge", open=False):
               levels = gr.Slider(label="Color levels", value=16, maximum=32, minimum=2, step=1)
@@ -612,7 +624,7 @@ with gr.Blocks(css=css, js=js) as demo:
               mouse = gr.Textbox(elem_id="mouse", value="""[]""", interactive=False)
               mouse.input(fn=draw_mask, show_progress="minimal", inputs=[levels, tolerance, mouse, output_mask], outputs=[output_mask])
               apply.click(fn=apply_mask, inputs=[output_mask, bsize], outputs=[output_mask, depth_file])
-            submit = gr.Button("Submit")
+              reset.click(fn=reset_mask, inputs=None, outputs=[output_mask, depth_file])
         with gr.Column():
             model_type = gr.Dropdown([("small", "vits"), ("base", "vitb"), ("large", "vitl")], type="value", value="vits", label='Model Type')
             processed_video = gr.Video(label="Output Video", format="mp4")