Update app.py

app.py

@@ -58,7 +58,7 @@ def predict_depth(model, image):
     return model(image)["depth"]
 
 #@spaces.GPU
-def make_video(video_path, outdir='./vis_video_depth', encoder='vits'):
+def make_video(video_path, outdir='./vis_video_depth', encoder='vits', remove_bg=False):
     if encoder not in ["vitl","vitb","vits"]:
         encoder = "vits"
 
@@ -196,7 +196,11 @@ def make_video(video_path, outdir='./vis_video_depth', encoder='vits'):
             masks.append(f"f{count}_mask.png")
             count += 1
 
-        final_vid = create_video(depth_frames, frame_rate, "depth")
+        if remove_bg == True:
+            final_vid = create_video(orig_frames, frame_rate, "orig")
+        else:
+            final_vid = create_video(depth_frames, frame_rate, "depth")
+            
         final_zip = zip_files(orig_frames, depth_frames)
         raw_video.release()
         # out.release()
@@ -935,6 +939,7 @@ with gr.Blocks(css=css, js=js) as demo:
 
         with gr.Column():
             model_type = gr.Dropdown([("small", "vits"), ("base", "vitb"), ("large", "vitl")], type="value", value="vits", label='Model Type')
+            remove_bg = gr.Checkbox(label="Remove background")
             processed_video = gr.Video(label="Output Video", format="mp4", interactive=False)
             processed_zip = gr.File(label="Output Archive", interactive=False)
             result = gr.Model3D(label="3D Mesh", clear_color=[0.5, 0.5, 0.5, 0.0], camera_position=[0, 90, 0], zoom_speed=2.0, pan_speed=2.0, interactive=True, elem_id="model3D") #, display_mode="point_cloud"
@@ -1271,7 +1276,7 @@ with gr.Blocks(css=css, js=js) as demo:
             render = gr.Button("Render")
             input_json.input(show_json, inputs=[input_json], outputs=[processed_video, processed_zip, output_frame, output_mask, output_depth, coords])
     
-    def on_submit(uploaded_video,model_type,coordinates):
+    def on_submit(uploaded_video,model_type,remove_bg,coordinates):
         global locations
         locations = []
         avg = [0, 0]
@@ -1311,11 +1316,11 @@ with gr.Blocks(css=css, js=js) as demo:
         print(locations)
             
         # Process the video and get the path of the output video
-        output_video_path = make_video(uploaded_video,encoder=model_type)
+        output_video_path = make_video(uploaded_video,encoder=model_type,remove_bg)
 
         return output_video_path + (json.dumps(locations),)
 
-    submit.click(on_submit, inputs=[input_video, model_type, coords], outputs=[processed_video, processed_zip, output_frame, output_mask, output_depth, coords])
+    submit.click(on_submit, inputs=[input_video, model_type, remove_bg, coords], outputs=[processed_video, processed_zip, output_frame, output_mask, output_depth, coords])
     render.click(None, inputs=[coords, mesh_order, bgcolor, output_frame, output_mask, selected, output_depth], outputs=None, js=load_model)
     render.click(partial(get_mesh), inputs=[output_frame, output_mask, blur_in, load_all], outputs=[result, result_file, mesh_order])