Update app.py

app.py

@@ -21,6 +21,8 @@ from depth_anything.util.transform import Resize, NormalizeImage, PrepareForNet
 from moviepy.editor import *
 
 frame_selected = 0
+frames = []
+depths = []
 masks = []
 locations = []
 mesh = []
@@ -170,8 +172,12 @@ def make_video(video_path, outdir='./vis_video_depth', encoder='vits'):
 
         global frame_selected
         global masks
+        global depths
+        global frames
+        frames = orig_frames
+        depths = depth_frames
         masks = orig_frames
-        return final_vid, final_zip, np.concatenate((orig_frames, depth_frames), axis=0), masks[frame_selected] #output_path
+        return final_vid, final_zip, orig_frames, masks[frame_selected] #output_path
 
 def depth_edges_mask(depth):
     """Returns a mask of edges in the depth map.
@@ -319,9 +325,13 @@ def select_frame(v, evt: gr.SelectData):
     #print(v)
     return v, frame_selected
 
-def align_rows(evt: gr.EventData):
-    global masks
-    return gr.Gallery(columns=int(len(masks)))
+def switch_rows(v):
+    global frames
+    global depths
+    if v == True:
+        return depths
+    else:
+        return frames
 
 
 css = """
@@ -370,6 +380,8 @@ with gr.Blocks(css=css) as demo:
             input_video = gr.Video(label="Input Video", format="mp4")
             input_url.change(fn=loadurl, inputs=[input_url], outputs=[input_video])
             output_frame = gr.Gallery(label="Frames", type='numpy', preview=True, columns=6)
+            output_switch = gr.Checkbox(label="Show depths")
+            output_switch.change(fn=switch_rows, inputs=[output_switch], outputs=[output_frame])
             output_frame.change(fn=align_rows, inputs=None, outputs=[output_frame])
             output_mask = gr.ImageEditor(interactive=True, transforms=(None,), eraser=gr.Eraser(), brush=gr.Brush(colors=['black', 'darkgray', 'gray', 'lightgray', 'white']), layers=True)
             submit = gr.Button("Submit")