prs-eth/rollingdepth
59
1# Copyright 2024 Anton Obukhov, ETH Zurich. All rights reserved.2#3# Licensed under the Apache License, Version 2.0 (the "License");4# you may not use this file except in compliance with the License.5# You may obtain a copy of the License at6#7# http://www.apache.org/licenses/LICENSE-2.08#9# Unless required by applicable law or agreed to in writing, software10# distributed under the License is distributed on an "AS IS" BASIS,11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.12# See the License for the specific language governing permissions and13# limitations under the License.14# --------------------------------------------------------------------------15# If you find this code useful, we kindly ask you to cite our paper in your work.16# Please find bibtex at: https://github.com/prs-eth/Marigold#-citation17# More information about the method can be found at https://marigoldmonodepth.github.io18# --------------------------------------------------------------------------19 20import functools21import os22import sys23import tempfile24 25import av26import numpy as np27 28import spaces29import gradio as gr30import torch as torch31import einops32 33from huggingface_hub import login34 35from colorize import colorize_depth_multi_thread36from video_io import get_video_fps, write_video_from_numpy37 38VERBOSE = False39MAX_FRAMES = 10040 41 42def process(pipe, device, path_input):43 print(f"Processing {path_input}")44 45 path_output_dir = tempfile.mkdtemp()46 os.makedirs(path_output_dir, exist_ok=True)47 48 name_base = os.path.splitext(os.path.basename(path_input))[0]49 path_out_in = os.path.join(path_output_dir, f"{name_base}_depth_input.mp4")50 path_out_vis = os.path.join(path_output_dir, f"{name_base}_depth_colored.mp4")51 52 output_fps = int(get_video_fps(path_input))53 54 container = av.open(path_input)55 stream = container.streams.video[0]56 fps = float(stream.average_rate)57 duration_sec = float(stream.duration * stream.time_base) if stream.duration else 058 total_frames = int(duration_sec * fps)59 if total_frames > MAX_FRAMES:60 gr.Warning(61 f"Only the first {MAX_FRAMES} frames (~{MAX_FRAMES / fps:.1f} sec.) will be processed for demonstration; "62 f"use the code from GitHub for full processing"63 )64 65 generator = torch.Generator(device=device)66 generator.manual_seed(2024)67 68 pipe_out: RollingDepthOutput = pipe(69 # input setting70 input_video_path=path_input,71 start_frame=0,72 frame_count=min(MAX_FRAMES, total_frames), # 0 = all73 processing_res=768,74 # infer setting75 dilations=[1, 25],76 cap_dilation=True,77 snippet_lengths=[3],78 init_infer_steps=[1],79 strides=[1],80 coalign_kwargs=None,81 refine_step=0, # 0 = off82 max_vae_bs=8, # batch size for encoder/decoder83 # other settings84 generator=generator,85 verbose=VERBOSE,86 # output settings87 restore_res=False,88 unload_snippet=False,89 )90 91 depth_pred = pipe_out.depth_pred # [N 1 H W]92 93 # Colorize results94 cmap = "Spectral_r"95 colored_np = colorize_depth_multi_thread(96 depth=depth_pred.numpy(),97 valid_mask=None,98 chunk_size=4,99 num_threads=4,100 color_map=cmap,101 verbose=VERBOSE,102 ) # [n h w 3], in [0, 255]103 104 write_video_from_numpy(105 frames=colored_np,106 output_path=path_out_vis,107 fps=output_fps,108 crf=23,109 preset="medium",110 verbose=VERBOSE,111 )112 113 # Save rgb114 rgb = (pipe_out.input_rgb.numpy() * 255).astype(np.uint8) # [N 3 H W]115 rgb = einops.rearrange(rgb, "n c h w -> n h w c")116 write_video_from_numpy(117 frames=rgb,118 output_path=path_out_in,119 fps=output_fps,120 crf=23,121 preset="medium",122 verbose=VERBOSE,123 )124 125 return path_out_in, path_out_vis126 127 128def run_demo_server(pipe, device):129 process_pipe = spaces.GPU(functools.partial(process, pipe, device), duration=120)130 os.environ["GRADIO_ALLOW_FLAGGING"] = "never"131 132 with gr.Blocks(133 analytics_enabled=False,134 title="RollingDepth",135 css="""136 h1 {137 text-align: center;138 display: block;139 }140 h2 {141 text-align: center;142 display: block;143 }144 h3 {145 text-align: center;146 display: block;147 }148 """,149 ) as demo:150 gr.HTML(151 """152 <h1>🛹 RollingDepth 🛹: Video Depth without Video Models</h1>153 <div style="text-align: center; margin-top: 20px;">154 <a title="Website" href="https://rollingdepth.github.io" target="_blank" rel="noopener noreferrer" style="display: inline-block; margin-right: 4px;">155 <img src="https://www.obukhov.ai/img/badges/badge-website.svg" alt="Website Badge">156 </a>157 <a title="arXiv" href="https://arxiv.org/abs/2411.19189" target="_blank" rel="noopener noreferrer" style="display: inline-block; margin-right: 4px;">158 <img src="https://www.obukhov.ai/img/badges/badge-pdf.svg" alt="arXiv Badge">159 </a>160 <a title="GitHub" href="https://github.com/prs-eth/rollingdepth" target="_blank" rel="noopener noreferrer" style="display: inline-block; margin-right: 4px;">161 <img src="https://img.shields.io/github/stars/prs-eth/rollingdepth?label=GitHub%20%E2%98%85&logo=github&color=C8C" alt="GitHub Stars Badge">162 </a>163 <a title="Social" href="https://twitter.com/antonobukhov1" target="_blank" rel="noopener noreferrer" style="display: inline-block; margin-right: 4px;">164 <img src="https://www.obukhov.ai/img/badges/badge-social.svg" alt="social">165 </a>166 </div>167 <p style="margin-top: 20px; text-align: justify;">168 RollingDepth is the state-of-the-art depth estimator for videos in the wild. Upload your video into the 169 <b>left</b> pane, or click any of the <b>examples</b> below. The result preview will be computed and 170 appear in the <b>right</b> panes. For full functionality, use the code on GitHub. 171 <b>TIP:</b> When running out of GPU time, fork the demo.172 </p>173 """174 )175 176 with gr.Row(equal_height=True):177 with gr.Column(scale=1):178 input_video = gr.Video(label="Input Video")179 with gr.Column(scale=2):180 with gr.Row(equal_height=True):181 output_video_1 = gr.Video(182 label="Preprocessed video",183 interactive=False,184 autoplay=True,185 loop=True,186 show_share_button=True,187 scale=5,188 )189 output_video_2 = gr.Video(190 label="Generated Depth Video",191 interactive=False,192 autoplay=True,193 loop=True,194 show_share_button=True,195 scale=5,196 )197 198 with gr.Row(equal_height=True):199 with gr.Column(scale=1):200 with gr.Row(equal_height=False):201 generate_btn = gr.Button("Generate")202 with gr.Column(scale=2):203 pass204 205 gr.Examples(206 examples=[207 ["files/gokart.mp4"],208 ["files/horse.mp4"],209 ["files/walking.mp4"],210 ],211 inputs=[input_video],212 outputs=[output_video_1, output_video_2],213 fn=process_pipe,214 cache_examples=True,215 cache_mode="eager",216 )217 218 generate_btn.click(219 fn=process_pipe,220 inputs=[input_video],221 outputs=[output_video_1, output_video_2],222 )223 224 demo.queue(225 api_open=False,226 ).launch(227 server_name="0.0.0.0",228 server_port=7860,229 )230 231 232def main():233 os.system("pip freeze")234 os.system("pip uninstall -y diffusers")235 os.system("pip install rollingdepth_src/diffusers")236 os.system("pip freeze")237 238 if "HF_TOKEN_LOGIN" in os.environ:239 login(token=os.environ["HF_TOKEN_LOGIN"])240 241 if torch.cuda.is_available():242 device = torch.device("cuda")243 elif torch.backends.mps.is_available():244 device = torch.device("mps")245 else:246 device = torch.device("cpu")247 248 sys.path.append(os.path.join(os.path.dirname(__file__), "rollingdepth_src"))249 from rollingdepth import RollingDepthOutput, RollingDepthPipeline250 251 pipe: RollingDepthPipeline = RollingDepthPipeline.from_pretrained(252 "prs-eth/rollingdepth-v1-0",253 torch_dtype=torch.float16,254 )255 pipe.set_progress_bar_config(disable=True)256 257 try:258 import xformers259 260 pipe.enable_xformers_memory_efficient_attention()261 except:262 pass # run without xformers263 264 pipe = pipe.to(device)265 266 run_demo_server(pipe, device)267 268 269if __name__ == "__main__":270 main()271