Update app.py

app.py

@@ -166,8 +166,9 @@ def make_video(video_path, outdir='./vis_video_depth', encoder='vits'):
         raw_video.release()
         # out.release()
         cv2.destroyAllWindows()
-        
-        return final_vid, final_zip, orig_frames, depth_frames #output_path
+
+        global frame_selected
+        return final_vid, final_zip, orig_frames, depth_frames, depth_frames[frame_selected] #output_path
 
 def depth_edges_mask(depth):
     """Returns a mask of edges in the depth map.
@@ -300,11 +301,11 @@ def blur_image(image, depth, blur_data):
 def loadurl(url):
     return url
 
-def select_frame(evt: gr.SelectData):
+def select_frame(v, evt: gr.SelectData):
     global frame_selected
     if evt.index != frame_selected:
         frame_selected = evt.index
-    return gr.Gallery(selected_index=evt.index, preview=True), frame_selected
+    return gr.Gallery(selected_index=evt.index, preview=True), gr.ImageEditor(background=v), frame_selected
 
 
 css = """
@@ -354,6 +355,7 @@ with gr.Blocks(css=css) as demo:
             input_url.change(fn=loadurl, inputs=[input_url], outputs=[input_video])
             output_frame = gr.Gallery(label="Frame", type='numpy', preview=True, columns=8192)
             output_depth = gr.Gallery(label="Depth", type='numpy', preview=True, columns=8192, interactive=False)
+            output_mask = gr.ImageEditor(transforms=None, eraser=True, brush=True, layers=True)
             submit = gr.Button("Submit")
         with gr.Column():
             model_type = gr.Dropdown([("small", "vits"), ("base", "vitb"), ("large", "vitl")], type="value", value="vits", label='Model Type')
@@ -564,8 +566,8 @@ with gr.Blocks(css=css) as demo:
       <a id='move' href='#'>move</a>    <a id='rotate' href='#'>rotate</a>    <a id='scale' href='#'>scale</a>
                 </pre>""")
                 selected = gr.Number(elem_id="fnum", value=0, minimum=0, maximum=256, interactive=False)
-                output_frame.select(fn=select_frame, inputs=None, outputs=[output_depth, selected], show_progress='hidden')
-                output_depth.select(fn=select_frame, inputs=None, outputs=[output_frame, selected], show_progress='hidden')
+                output_frame.select(fn=select_frame, inputs=output_depth, outputs=[output_depth, output_mask, selected], show_progress='hidden')
+                output_depth.select(fn=select_frame, inputs=output_depth, outputs=[output_frame, output_mask, selected], show_progress='hidden')
                 example_coords = """[
                   {"latLng": { "lat": 50.07379596793083, "lng": 14.437146122950555 } },
                   {"latLng": { "lat": 50.073799567020004, "lng": 14.437146774240507 } },
@@ -719,14 +721,14 @@ with gr.Blocks(css=css) as demo:
 
         return output_video_path + (locations,)
 
-    submit.click(on_submit, inputs=[input_video, model_type, coords], outputs=[processed_video, processed_zip, output_frame, output_depth, coords])
+    submit.click(on_submit, inputs=[input_video, model_type, coords], outputs=[processed_video, processed_zip, output_frame, output_depth, output_mask, coords])
     render.click(partial(get_mesh), inputs=[output_frame, output_depth, blur_in, load_all], outputs=[result])
 
     example_files = os.listdir('examples')
     example_files.sort()
     example_files = [os.path.join('examples', filename) for filename in example_files]
     
-    examples = gr.Examples(examples=example_files, inputs=[input_video], outputs=[processed_video, processed_zip, output_frame, output_depth, coords], fn=on_submit, cache_examples=True)
+    examples = gr.Examples(examples=example_files, inputs=[input_video], outputs=[processed_video, processed_zip, output_frame, output_depth, output_mask, coords], fn=on_submit, cache_examples=True)
     
 
 if __name__ == '__main__':