Buckets:
| # Copyright (2025) Bytedance Ltd. and/or its affiliates | |
| # Licensed under the Apache License, Version 2.0 (the "License"); | |
| # you may not use this file except in compliance with the License. | |
| # You may obtain a copy of the License at | |
| # http://www.apache.org/licenses/LICENSE-2.0 | |
| # Unless required by applicable law or agreed to in writing, software | |
| # distributed under the License is distributed on an "AS IS" BASIS, | |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | |
| # See the License for the specific language governing permissions and | |
| # limitations under the License. | |
| import argparse | |
| import numpy as np | |
| import os | |
| import torch | |
| from video_depth_anything.video_depth import VideoDepthAnything | |
| from utils.dc_utils import read_video_frames, save_video | |
| if __name__ == '__main__': | |
| parser = argparse.ArgumentParser(description='Video Depth Anything') | |
| parser.add_argument('--input_video', type=str, default='./assets/example_videos/davis_rollercoaster.mp4') | |
| parser.add_argument('--output_dir', type=str, default='./outputs') | |
| parser.add_argument('--input_size', type=int, default=518) | |
| parser.add_argument('--max_res', type=int, default=1280) | |
| parser.add_argument('--encoder', type=str, default='vitl', choices=['vits', 'vitb', 'vitl']) | |
| parser.add_argument('--max_len', type=int, default=-1, help='maximum length of the input video, -1 means no limit') | |
| parser.add_argument('--target_fps', type=int, default=-1, help='target fps of the input video, -1 means the original fps') | |
| parser.add_argument('--metric', action='store_true', help='use metric model') | |
| parser.add_argument('--fp32', action='store_true', help='model infer with torch.float32, default is torch.float16') | |
| parser.add_argument('--grayscale', action='store_true', help='do not apply colorful palette') | |
| parser.add_argument('--save_npz', action='store_true', help='save depths as npz') | |
| parser.add_argument('--save_exr', action='store_true', help='save depths as exr') | |
| parser.add_argument('--focal-length-x', default=470.4, type=float, | |
| help='Focal length along the x-axis.') | |
| parser.add_argument('--focal-length-y', default=470.4, type=float, | |
| help='Focal length along the y-axis.') | |
| args = parser.parse_args() | |
| DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu' | |
| model_configs = { | |
| 'vits': {'encoder': 'vits', 'features': 64, 'out_channels': [48, 96, 192, 384]}, | |
| 'vitb': {'encoder': 'vitb', 'features': 128, 'out_channels': [96, 192, 384, 768]}, | |
| 'vitl': {'encoder': 'vitl', 'features': 256, 'out_channels': [256, 512, 1024, 1024]}, | |
| } | |
| checkpoint_name = 'metric_video_depth_anything' if args.metric else 'video_depth_anything' | |
| video_depth_anything = VideoDepthAnything(**model_configs[args.encoder], metric=args.metric) | |
| video_depth_anything.load_state_dict(torch.load(f'./checkpoints/{checkpoint_name}_{args.encoder}.pth', map_location='cpu'), strict=True) | |
| video_depth_anything = video_depth_anything.to(DEVICE).eval() | |
| frames, target_fps = read_video_frames(args.input_video, args.max_len, args.target_fps, args.max_res) | |
| depths, fps = video_depth_anything.infer_video_depth(frames, target_fps, input_size=args.input_size, device=DEVICE, fp32=args.fp32) | |
| video_name = os.path.basename(args.input_video) | |
| os.makedirs(args.output_dir, exist_ok=True) | |
| processed_video_path = os.path.join(args.output_dir, os.path.splitext(video_name)[0]+'_src.mp4') | |
| depth_vis_path = os.path.join(args.output_dir, os.path.splitext(video_name)[0]+'_vis.mp4') | |
| save_video(frames, processed_video_path, fps=fps) | |
| save_video(depths, depth_vis_path, fps=fps, is_depths=True, grayscale=args.grayscale) | |
| if args.save_npz: | |
| depth_npz_path = os.path.join(args.output_dir, os.path.splitext(video_name)[0]+'_depths.npz') | |
| np.savez_compressed(depth_npz_path, depths=depths) | |
| if args.save_exr: | |
| depth_exr_dir = os.path.join(args.output_dir, os.path.splitext(video_name)[0]+'_depths_exr') | |
| os.makedirs(depth_exr_dir, exist_ok=True) | |
| import OpenEXR | |
| import Imath | |
| for i, depth in enumerate(depths): | |
| output_exr = f"{depth_exr_dir}/frame_{i:05d}.exr" | |
| header = OpenEXR.Header(depth.shape[1], depth.shape[0]) | |
| header["channels"] = { | |
| "Z": Imath.Channel(Imath.PixelType(Imath.PixelType.FLOAT)) | |
| } | |
| exr_file = OpenEXR.OutputFile(output_exr, header) | |
| exr_file.writePixels({"Z": depth.tobytes()}) | |
| exr_file.close() | |
| if args.metric: | |
| import open3d as o3d | |
| width, height = depths[0].shape[-1], depths[0].shape[-2] | |
| x, y = np.meshgrid(np.arange(width), np.arange(height)) | |
| x = (x - width / 2) / args.focal_length_x | |
| y = (y - height / 2) / args.focal_length_y | |
| for i, (color_image, depth) in enumerate(zip(frames, depths)): | |
| z = np.array(depth) | |
| points = np.stack((np.multiply(x, z), np.multiply(y, z), z), axis=-1).reshape(-1, 3) | |
| colors = np.array(color_image).reshape(-1, 3) / 255.0 | |
| pcd = o3d.geometry.PointCloud() | |
| pcd.points = o3d.utility.Vector3dVector(points) | |
| pcd.colors = o3d.utility.Vector3dVector(colors) | |
| o3d.io.write_point_cloud(os.path.join(args.output_dir, 'point' + str(i).zfill(4) + '.ply'), pcd) | |
Xet Storage Details
- Size:
- 5.42 kB
- Xet hash:
- c43312257dfba3a09c7c2ec92a9bc645bb66e87b1ee27a4c5b89927ca70fc1ba
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.