Spaces:
Running
on
Zero
Running
on
Zero
File size: 3,481 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 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 |
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.
# # Data loading based on https://github.com/NVIDIA/flownet2-pytorch
import os
import numpy as np
import os.path as osp
from PIL import Image
from tqdm import tqdm
import csv
import imageio
import cv2
import json
import glob
import shutil
from eval_utils import gen_json, get_sorted_files
TAG_FLOAT = 202021.25
TAG_CHAR = "PIEH"
def depth_read(filename):
"""Read depth data from file, return as numpy array."""
f = open(filename, "rb")
check = np.fromfile(f, dtype=np.float32, count=1)[0]
assert (
check == TAG_FLOAT
), " depth_read:: Wrong tag in flow file (should be: {0}, is: {1}). Big-endian machine? ".format(
TAG_FLOAT, check
)
width = np.fromfile(f, dtype=np.int32, count=1)[0]
height = np.fromfile(f, dtype=np.int32, count=1)[0]
size = width * height
assert (
width > 0 and height > 0 and size > 1 and size < 100000000
), " depth_read:: Wrong input size (width = {0}, height = {1}).".format(
width, height
)
depth = np.fromfile(f, dtype=np.float32, count=-1).reshape((height, width))
return depth
def extract_sintel(
root,
depth_root,
sample_len=-1,
datatset_name="",
saved_dir="",
):
scenes_names = os.listdir(root)
all_samples = []
for i, seq_name in enumerate(tqdm(scenes_names)):
all_img_names = get_sorted_files(
os.path.join(root, seq_name), suffix=".png")
seq_len = len(all_img_names)
step = sample_len if sample_len > 0 else seq_len
for ref_idx in range(0, seq_len, step):
print(f"Progress: {seq_name}, {ref_idx // step} / {seq_len // step}")
if (ref_idx + step) <= seq_len:
ref_e = ref_idx + step
else:
continue
for idx in range(ref_idx, ref_e):
im_path = osp.join(
root, seq_name, all_img_names[idx]
)
depth_path = osp.join(
depth_root, seq_name, all_img_names[idx][:-3] + "dpt"
)
out_img_path = osp.join(
saved_dir, datatset_name,'clean', seq_name, all_img_names[idx]
)
out_depth_path = osp.join(
saved_dir, datatset_name,'depth', seq_name, all_img_names[idx][:-3] + "png"
)
depth = depth_read(depth_path)
img = np.array(Image.open(im_path))
os.makedirs(osp.dirname(out_img_path), exist_ok=True)
os.makedirs(osp.dirname(out_depth_path), exist_ok=True)
cv2.imwrite(
out_img_path,
img,
)
cv2.imwrite(
out_depth_path,
depth.astype(np.uint16)
)
gen_json(
root_path=osp.join(saved_dir, datatset_name), dataset=datatset_name,
start_id=0,end_id=100,step=1,
save_path=osp.join(saved_dir, datatset_name, "sintel_video.json"),)
if __name__ == "__main__":
extract_sintel(
root="path/to/training/clean",
depth_root="path/to/depth",
saved_dir="./benchmark/datasets/",
sample_len=-1,
datatset_name="sintel",
) |