long_app.py

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)