Spaces:
Running
on
Zero
Running
on
Zero
File size: 2,620 Bytes
5a0778e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 |
import argparse
import os
import cv2
import json
import torch
from tqdm import tqdm
import numpy as np
from video_depth_anything.video_depth import VideoDepthAnything
from utils.dc_utils import read_video_frames
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument('--infer_path', type=str, default='')
parser.add_argument('--json_file', type=str, default='')
parser.add_argument('--datasets', type=str, nargs='+', default=['scannet', 'nyuv2'])
parser.add_argument('--input_size', type=int, default=518)
parser.add_argument('--encoder', type=str, default='vitl', choices=['vits', 'vitl'])
args = parser.parse_args()
for dataset in args.datasets:
with open(args.json_file, 'r') as fs:
path_json = json.load(fs)
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]},
}
video_depth_anything = VideoDepthAnything(**model_configs[args.encoder])
video_depth_anything.load_state_dict(torch.load(f'./checkpoints/video_depth_anything_{args.encoder}.pth', map_location='cpu'), strict=True)
video_depth_anything = video_depth_anything.to(DEVICE).eval()
json_data = path_json[dataset]
root_path = os.path.dirname(args.json_file)
for data in tqdm(json_data):
for key in data.keys():
value = data[key]
infer_paths = []
videos = []
for images in value:
image_path = os.path.join(root_path, images['image'])
infer_path = (args.infer_path + '/'+ dataset + '/' + images['image']).replace('.jpg', '.npy').replace('.png', '.npy')
infer_paths.append(infer_path)
img = cv2.imread(image_path)
videos.append(img)
videos = np.stack(videos, axis=0)
target_fps=1
depths, fps = video_depth_anything.infer_video_depth(videos, target_fps, input_size=args.input_size, device=DEVICE, fp32=True)
for i in range(len(infer_paths)):
infer_path = infer_paths[i]
os.makedirs(os.path.dirname(infer_path), exist_ok=True)
depth = depths[i]
np.save(infer_path, depth)
|