Spaces:
Running
on
Zero
Running
on
Zero
| import os | |
| import numpy as np | |
| import os.path as osp | |
| from PIL import Image | |
| from tqdm import tqdm | |
| import csv | |
| import imageio | |
| def _read_image(img_rel_path) -> np.ndarray: | |
| image_to_read = img_rel_path | |
| image = Image.open(image_to_read) | |
| image = np.asarray(image) | |
| return image | |
| def depth_read(filename): | |
| depth_in = _read_image(filename) | |
| depth_decoded = depth_in / 1000.0 | |
| return depth_decoded | |
| def extract_nyu( | |
| root, | |
| depth_root, | |
| csv_save_path="", | |
| datatset_name="", | |
| filename_ls_path="", | |
| saved_rgb_dir="", | |
| saved_disp_dir="", | |
| ): | |
| with open(filename_ls_path, "r") as f: | |
| filenames = [s.split() for s in f.readlines()] | |
| all_samples = [] | |
| for i, pair_names in enumerate(tqdm(filenames)): | |
| img_name = pair_names[0] | |
| filled_depth_name = pair_names[2] | |
| im_path = osp.join(root, img_name) | |
| depth_path = osp.join(depth_root, filled_depth_name) | |
| depth = depth_read(depth_path) | |
| disp = depth | |
| video_depths = [disp] | |
| video_imgs = [np.array(Image.open(im_path))] | |
| disp_video = np.array(video_depths)[:, None] | |
| img_video = np.array(video_imgs)[..., 0:3] | |
| disp_video = disp_video[:, :, 45:471, 41:601] | |
| img_video = img_video[:, 45:471, 41:601, :] | |
| data_root = saved_rgb_dir + datatset_name | |
| disp_root = saved_disp_dir + datatset_name | |
| os.makedirs(data_root, exist_ok=True) | |
| os.makedirs(disp_root, exist_ok=True) | |
| img_video_dir = data_root | |
| disp_video_dir = disp_root | |
| img_video_path = os.path.join(img_video_dir, f"{img_name[:-4]}_rgb_left.mp4") | |
| disp_video_path = os.path.join(disp_video_dir, f"{img_name[:-4]}_disparity.npz") | |
| dir_name = os.path.dirname(img_video_path) | |
| os.makedirs(dir_name, exist_ok=True) | |
| dir_name = os.path.dirname(disp_video_path) | |
| os.makedirs(dir_name, exist_ok=True) | |
| imageio.mimsave( | |
| img_video_path, img_video, fps=15, quality=10, macro_block_size=1 | |
| ) | |
| np.savez(disp_video_path, disparity=disp_video) | |
| sample = {} | |
| sample["filepath_left"] = os.path.join( | |
| f"{datatset_name}/{img_name[:-4]}_rgb_left.mp4" | |
| ) | |
| sample["filepath_disparity"] = os.path.join( | |
| f"{datatset_name}/{img_name[:-4]}_disparity.npz" | |
| ) | |
| all_samples.append(sample) | |
| filename_ = csv_save_path | |
| os.makedirs(os.path.dirname(filename_), exist_ok=True) | |
| fields = ["filepath_left", "filepath_disparity"] | |
| with open(filename_, "w") as csvfile: | |
| writer = csv.DictWriter(csvfile, fieldnames=fields) | |
| writer.writeheader() | |
| writer.writerows(all_samples) | |
| print(f"{filename_} has been saved.") | |
| if __name__ == "__main__": | |
| extract_nyu( | |
| root="path/to/NYUv2/", | |
| depth_root="path/to/NYUv2/", | |
| filename_ls_path="path/to/NYUv2/filename_list_test.txt", | |
| saved_rgb_dir="./benchmark/datasets/", | |
| saved_disp_dir="./benchmark/datasets/", | |
| csv_save_path=f"./benchmark/datasets/NYUv2.csv", | |
| datatset_name="NYUv2", | |
| ) | |