# MIT License
# Copyright (c) 2024 Jiahao Shao
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.
import functools
import os
import zipfile
import tempfile
from io import BytesIO
import spaces
import gradio as gr
import numpy as np
import torch as torch
from PIL import Image
from tqdm import tqdm
import mediapy as media
from huggingface_hub import login
from chronodepth_pipeline import ChronoDepthPipeline
from gradio_patches.examples import Examples
default_seed = 2024
default_num_inference_steps = 5
default_num_frames = 10
default_window_size = 9
default_video_processing_resolution = 768
default_video_out_max_frames = 80
default_decode_chunk_size = 10
def process_video(
pipe,
path_input,
num_inference_steps=default_num_inference_steps,
num_frames=default_num_frames,
window_size=default_window_size,
out_max_frames=default_video_out_max_frames,
progress=gr.Progress(),
):
if path_input is None:
raise gr.Error(
"Missing video in the first pane: upload a file or use one from the gallery below."
)
name_base, name_ext = os.path.splitext(os.path.basename(path_input))
print(f"Processing video {name_base}{name_ext}")
path_output_dir = tempfile.mkdtemp()
path_out_vis = os.path.join(path_output_dir, f"{name_base}_depth_colored.mp4")
path_out_16bit = os.path.join(path_output_dir, f"{name_base}_depth_16bit.zip")
generator = torch.Generator(device=pipe.device).manual_seed(default_seed)
import time
start_time = time.time()
zipf = None
try:
if window_size is None or window_size == num_frames:
inpaint_inference = False
else:
inpaint_inference = True
data_ls = []
video_data = media.read_video(path_input)
video_length = len(video_data)
fps = video_data.metadata.fps
duration_sec = video_length / fps
out_duration_sec = out_max_frames / fps
if duration_sec > out_duration_sec:
gr.Warning(
f"Only the first ~{int(out_duration_sec)} seconds will be processed; "
f"use alternative setups such as ChronoDepth on github for full processing"
)
video_length = out_max_frames
for i in tqdm(range(video_length-num_frames+1)):
is_first_clip = i == 0
is_last_clip = i == video_length - num_frames
is_new_clip = (
(inpaint_inference and i % window_size == 0)
or (inpaint_inference == False and i % num_frames == 0)
)
if is_first_clip or is_last_clip or is_new_clip:
data_ls.append(np.array(video_data[i: i+num_frames])) # [t, H, W, 3]
zipf = zipfile.ZipFile(path_out_16bit, "w", zipfile.ZIP_DEFLATED)
depth_colored_pred = []
depth_pred = []
# -------------------- Inference and saving --------------------
with torch.no_grad():
for iter, batch in enumerate(tqdm(data_ls)):
rgb_int = batch
input_images = [Image.fromarray(rgb_int[i]) for i in range(num_frames)]
# Predict depth
if iter == 0: # First clip
pipe_out = pipe(
input_images,
num_frames=len(input_images),
num_inference_steps=num_inference_steps,
decode_chunk_size=default_decode_chunk_size,
motion_bucket_id=127,
fps=7,
noise_aug_strength=0.0,
generator=generator,
)
elif inpaint_inference and (iter == len(data_ls) - 1): # temporal inpaint inference for last clip
last_window_size = window_size if video_length%window_size == 0 else video_length%window_size
pipe_out = pipe(
input_images,
num_frames=num_frames,
num_inference_steps=num_inference_steps,
decode_chunk_size=default_decode_chunk_size,
motion_bucket_id=127,
fps=7,
noise_aug_strength=0.0,
generator=generator,
depth_pred_last=depth_frames_pred_ts[last_window_size:],
)
elif inpaint_inference and iter > 0: # temporal inpaint inference
pipe_out = pipe(
input_images,
num_frames=num_frames,
num_inference_steps=num_inference_steps,
decode_chunk_size=default_decode_chunk_size,
motion_bucket_id=127,
fps=7,
noise_aug_strength=0.0,
generator=generator,
depth_pred_last=depth_frames_pred_ts[window_size:],
)
else: # separate inference
pipe_out = pipe(
input_images,
num_frames=num_frames,
num_inference_steps=num_inference_steps,
decode_chunk_size=default_decode_chunk_size,
motion_bucket_id=127,
fps=7,
noise_aug_strength=0.0,
generator=generator,
)
depth_frames_pred = [pipe_out.depth_np[i] for i in range(num_frames)]
depth_frames_colored_pred = []
for i in range(num_frames):
depth_frame_colored_pred = np.array(pipe_out.depth_colored[i])
depth_frames_colored_pred.append(depth_frame_colored_pred)
depth_frames_colored_pred = np.stack(depth_frames_colored_pred, axis=0)
depth_frames_pred = np.stack(depth_frames_pred, axis=0)
depth_frames_pred_ts = torch.from_numpy(depth_frames_pred).to(pipe.device)
depth_frames_pred_ts = depth_frames_pred_ts * 2 - 1
if inpaint_inference == False:
if iter == len(data_ls) - 1:
last_window_size = num_frames if video_length%num_frames == 0 else video_length%num_frames
depth_colored_pred.append(depth_frames_colored_pred[-last_window_size:])
depth_pred.append(depth_frames_pred[-last_window_size:])
else:
depth_colored_pred.append(depth_frames_colored_pred)
depth_pred.append(depth_frames_pred)
else:
if iter == 0:
depth_colored_pred.append(depth_frames_colored_pred)
depth_pred.append(depth_frames_pred)
elif iter == len(data_ls) - 1:
depth_colored_pred.append(depth_frames_colored_pred[-last_window_size:])
depth_pred.append(depth_frames_pred[-last_window_size:])
else:
depth_colored_pred.append(depth_frames_colored_pred[-window_size:])
depth_pred.append(depth_frames_pred[-window_size:])
depth_colored_pred = np.concatenate(depth_colored_pred, axis=0)
depth_pred = np.concatenate(depth_pred, axis=0)
# -------------------- Save results --------------------
# Save images
for i in tqdm(range(len(depth_pred))):
archive_path = os.path.join(
f"{name_base}_depth_16bit", f"{i:05d}.png"
)
img_byte_arr = BytesIO()
depth_16bit = Image.fromarray((depth_pred[i] * 65535.0).astype(np.uint16))
depth_16bit.save(img_byte_arr, format="png")
img_byte_arr.seek(0)
zipf.writestr(archive_path, img_byte_arr.read())
# Export to video
media.write_video(path_out_vis, depth_colored_pred, fps=fps)
finally:
if zipf is not None:
zipf.close()
end_time = time.time()
print(f"Processing time: {end_time - start_time} seconds")
return (
path_out_vis,
[path_out_vis, path_out_16bit],
)
def run_demo_server(pipe):
process_pipe_video = spaces.GPU(
functools.partial(process_video, pipe), duration=220
)
os.environ["GRADIO_ALLOW_FLAGGING"] = "never"
with gr.Blocks(
analytics_enabled=False,
title="ChronoDepth Video Depth Estimation",
css="""
#download {
height: 118px;
}
.slider .inner {
width: 5px;
background: #FFF;
}
.viewport {
aspect-ratio: 4/3;
}
h1 {
text-align: center;
display: block;
}
h2 {
text-align: center;
display: block;
}
h3 {
text-align: center;
display: block;
}
""",
) as demo:
gr.Markdown(
"""
# ChronoDepth Video Depth Estimation
ChronoDepth is the state-of-the-art video depth estimator for videos in the wild.
Upload your video and have a try!
We set denoising steps to 5, number of frames for each video clip to 10, and overlap between clips to 1.
"""
)
with gr.Row():
with gr.Column():
video_input = gr.Video(
label="Input Video",
sources=["upload"],
)
with gr.Row():
video_submit_btn = gr.Button(
value="Compute Depth", variant="primary"
)
video_reset_btn = gr.Button(value="Reset")
with gr.Column():
video_output_video = gr.Video(
label="Output video depth (red-near, blue-far)",
interactive=False,
)
video_output_files = gr.Files(
label="Depth outputs",
elem_id="download",
interactive=False,
)
Examples(
fn=process_pipe_video,
examples=[
os.path.join("files", name)
for name in [
"sora_e2.mp4",
"sora_1758192960116785459.mp4",
]
],
inputs=[video_input],
outputs=[video_output_video, video_output_files],
cache_examples=True,
directory_name="examples_video",
)
video_submit_btn.click(
fn=process_pipe_video,
inputs=[video_input],
outputs=[video_output_video, video_output_files],
concurrency_limit=1,
)
video_reset_btn.click(
fn=lambda: (None, None, None),
inputs=[],
outputs=[video_input, video_output_video],
concurrency_limit=1,
)
demo.queue(
api_open=False,
).launch(
server_name="0.0.0.0",
server_port=7860,
)
def main():
CHECKPOINT = "jhshao/ChronoDepth"
if "HF_TOKEN_LOGIN" in os.environ:
login(token=os.environ["HF_TOKEN_LOGIN"])
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Running on device: {device}")
pipe = ChronoDepthPipeline.from_pretrained(CHECKPOINT)
try:
import xformers
pipe.enable_xformers_memory_efficient_attention()
except:
pass # run without xformers
pipe = pipe.to(device)
run_demo_server(pipe)
if __name__ == "__main__":
main()