Upload long_app.py

long_app.py

@@ -0,0 +1,207 @@
+import os
+import gc
+import numpy as np
+import torch
+import spaces
+import gradio as gr
+from moviepy.editor import VideoFileClip, concatenate_videoclips
+from video_depth_anything.video_depth import VideoDepthAnything
+from utils.dc_utils import read_video_frames, save_video
+from huggingface_hub import hf_hub_download
+
+examples = [
+    ['assets/example_videos/davis_rollercoaster.mp4', -1, -1, 1280],
+    ['assets/example_videos/Tokyo-Walk_rgb.mp4', -1, -1, 1280],
+    ['assets/example_videos/4158877-uhd_3840_2160_30fps_rgb.mp4', -1, -1, 1280],
+    ['assets/example_videos/4511004-uhd_3840_2160_24fps_rgb.mp4', -1, -1, 1280],
+    ['assets/example_videos/1753029-hd_1920_1080_30fps.mp4', -1, -1, 1280],
+    ['assets/example_videos/davis_burnout.mp4', -1, -1, 1280],
+    ['assets/example_videos/example_5473765-l.mp4', -1, -1, 1280],
+    ['assets/example_videos/Istanbul-26920.mp4', -1, -1, 1280],
+    ['assets/example_videos/obj_1.mp4', -1, -1, 1280],
+    ['assets/example_videos/sheep_cut1.mp4', -1, -1, 1280],
+]
+
+DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
+
+model_configs = {
+    'vits': {'encoder': 'vits', 'features': 64, 'out_channels': [48, 96, 192, 384]},
+    'vitl': {'encoder': 'vitl', 'features': 256, 'out_channels': [256, 512, 1024, 1024]},
+}
+
+encoder2name = {
+    'vits': 'Small',
+    'vitl': 'Large',
+}
+
+#encoder = 'vitl'
+encoder = 'vits'
+model_name = encoder2name[encoder]
+
+video_depth_anything = VideoDepthAnything(**model_configs[encoder])
+filepath = hf_hub_download(repo_id=f"depth-anything/Video-Depth-Anything-{model_name}", filename=f"video_depth_anything_{encoder}.pth", repo_type="model")
+video_depth_anything.load_state_dict(torch.load(filepath, map_location='cpu'))
+video_depth_anything = video_depth_anything.to(DEVICE).eval()
+
+title = "# Video Depth Anything"
+description = """Official demo for ​**Video Depth Anything**.
+Please refer to our [paper](https://arxiv.org/abs/2501.12375), [project page](https://videodepthanything.github.io/), and [github](https://github.com/DepthAnything/Video-Depth-Anything) for more details."""
+
+@spaces.GPU(duration=240)
+def infer_video_depth(
+    input_video: str,
+    max_len: int = -1,
+    target_fps: int = -1,
+    max_res: int = 1280,
+    grayscale: bool = False,
+    output_dir: str = './outputs',
+    input_size: int = 518,
+):
+    if not os.path.exists(output_dir):
+        os.makedirs(output_dir)
+
+    video_name = os.path.basename(input_video)
+    processed_video_path = os.path.join(output_dir, os.path.splitext(video_name)[0]+'_src.mp4')
+    depth_vis_path = os.path.join(output_dir, os.path.splitext(video_name)[0]+'_vis.mp4')
+
+    # Load the video
+    clip = VideoFileClip(input_video)
+    fps = clip.fps
+    total_frames = int(clip.duration * fps)
+
+    # Define the number of frames per segment
+    frames_per_segment = 45  # Adjust this value based on your GPU memory
+    segments = []
+    for start_frame in range(0, total_frames, frames_per_segment):
+        end_frame = min(start_frame + frames_per_segment, total_frames)
+        start_time = start_frame / fps
+        end_time = end_frame / fps
+        segment = clip.subclip(start_time, end_time)
+        segment_path = os.path.join(output_dir, f'segment_{start_frame}.mp4')
+        segment.write_videofile(segment_path, codec='libx264')
+        segments.append(segment_path)
+
+    # Save the processed video (concatenated segments)
+    processed_segments = [VideoFileClip(segment) for segment in segments]
+    final_processed_clip = concatenate_videoclips(processed_segments)
+    final_processed_clip.write_videofile(processed_video_path, codec='libx264')
+
+    # Process each segment
+    depth_segments = []
+    for segment in segments:
+        frames, target_fps = read_video_frames(segment, max_len, target_fps, max_res)
+        print("frame length", len(frames))
+        depths, fps = video_depth_anything.infer_video_depth(frames, target_fps, input_size=input_size, device=DEVICE)
+        depth_segment_path = os.path.join(output_dir, f'depth_{os.path.basename(segment)}')
+        save_video(depths, depth_segment_path, fps=fps, is_depths=True, grayscale=grayscale)
+        depth_segments.append(depth_segment_path)
+
+    # Merge depth segments
+    depth_clips = [VideoFileClip(depth_segment) for depth_segment in depth_segments]
+    final_depth_clip = concatenate_videoclips(depth_clips)
+    final_depth_clip.write_videofile(depth_vis_path, codec='libx264')
+
+    # Clean up
+    for segment in segments:
+        os.remove(segment)
+    for depth_segment in depth_segments:
+        os.remove(depth_segment)
+
+    gc.collect()
+    torch.cuda.empty_cache()
+
+    return [processed_video_path, depth_vis_path]
+
+def construct_demo():
+    with gr.Blocks(analytics_enabled=False) as demo:
+        gr.Markdown(title)
+        gr.Markdown(description)
+        gr.Markdown("### If you find this work useful, please help ⭐ the [$$Github Repo$$](https://github.com/DepthAnything/Video-Depth-Anything). Thanks for your attention!")
+
+        with gr.Row(equal_height=True):
+            with gr.Column(scale=1):
+                input_video = gr.Video(label="Input Video")
+
+            with gr.Column(scale=2):
+                with gr.Row(equal_height=True):
+                    processed_video = gr.Video(
+                        label="Preprocessed video",
+                        interactive=False,
+                        autoplay=True,
+                        loop=True,
+                        show_share_button=True,
+                        scale=5,
+                    )
+                    depth_vis_video = gr.Video(
+                        label="Generated Depth Video",
+                        interactive=False,
+                        autoplay=True,
+                        loop=True,
+                        show_share_button=True,
+                        scale=5,
+                    )
+
+        with gr.Row(equal_height=True):
+            with gr.Column(scale=1):
+                with gr.Row(equal_height=False):
+                    with gr.Accordion("Advanced Settings", open=False):
+                        max_len = gr.Slider(
+                            label="max process length",
+                            minimum=-1,
+                            maximum=1000,
+                            value=500,
+                            step=1,
+                        )
+                        target_fps = gr.Slider(
+                            label="target FPS",
+                            minimum=-1,
+                            maximum=30,
+                            value=15,
+                            step=1,
+                        )
+                        max_res = gr.Slider(
+                            label="max side resolution",
+                            minimum=480,
+                            maximum=1920,
+                            value=1280,
+                            step=1,
+                        )
+                        grayscale = gr.Checkbox(
+                            label="grayscale",
+                            value=False,
+                        )
+                    generate_btn = gr.Button("Generate")
+            with gr.Column(scale=2):
+                pass
+
+        gr.Examples(
+            examples=examples,
+            inputs=[
+                input_video,
+                max_len,
+                target_fps,
+                max_res
+            ],
+            outputs=[processed_video, depth_vis_video],
+            fn=infer_video_depth,
+            cache_examples="lazy",
+        )
+
+        generate_btn.click(
+            fn=infer_video_depth,
+            inputs=[
+                input_video,
+                max_len,
+                target_fps,
+                max_res,
+                grayscale
+            ],
+            outputs=[processed_video, depth_vis_video],
+        )
+
+    return demo
+
+if __name__ == "__main__":
+    demo = construct_demo()
+    demo.queue()
+    demo.launch(share=True)