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 imageio | |
| import csv | |
| def depth_read(filename): | |
| # loads depth map D from png file | |
| # and returns it as a numpy array | |
| depth_png = np.asarray(Image.open(filename)) | |
| # make sure we have a proper 16bit depth map here.. not 8bit! | |
| assert np.max(depth_png) > 255 | |
| depth = depth_png.astype(np.float64) / 5000.0 | |
| depth[depth_png == 0] = -1.0 | |
| return depth | |
| def extract_bonn( | |
| root, | |
| depth_root, | |
| sample_len=-1, | |
| csv_save_path="", | |
| datatset_name="", | |
| saved_rgb_dir="", | |
| saved_disp_dir="", | |
| start_frame=0, | |
| end_frame=110, | |
| ): | |
| scenes_names = os.listdir(depth_root) | |
| all_samples = [] | |
| for i, seq_name in enumerate(tqdm(scenes_names)): | |
| # load all images | |
| all_img_names = os.listdir(osp.join(depth_root, seq_name, "rgb")) | |
| all_img_names = [x for x in all_img_names if x.endswith(".png")] | |
| print(f"sequence frame number: {len(all_img_names)}") | |
| # for not zero padding image name | |
| all_img_names.sort() | |
| all_img_names = sorted(all_img_names, key=lambda x: int(x.split(".")[0][-4:])) | |
| all_img_names = all_img_names[start_frame:end_frame] | |
| all_depth_names = os.listdir(osp.join(depth_root, seq_name, "depth")) | |
| all_depth_names = [x for x in all_depth_names if x.endswith(".png")] | |
| print(f"sequence depth number: {len(all_depth_names)}") | |
| # for not zero padding image name | |
| all_depth_names.sort() | |
| all_depth_names = sorted( | |
| all_depth_names, key=lambda x: int(x.split(".")[0][-4:]) | |
| ) | |
| all_depth_names = all_depth_names[start_frame:end_frame] | |
| 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 + 1} / {seq_len//step}") | |
| video_imgs = [] | |
| video_depths = [] | |
| if (ref_idx + step) <= seq_len: | |
| ref_e = ref_idx + step | |
| else: | |
| continue | |
| # for idx in range(ref_idx, ref_idx + step): | |
| for idx in range(ref_idx, ref_e): | |
| im_path = osp.join(root, seq_name, "rgb", all_img_names[idx]) | |
| depth_path = osp.join( | |
| depth_root, seq_name, "depth", all_depth_names[idx] | |
| ) | |
| depth = depth_read(depth_path) | |
| disp = depth | |
| video_depths.append(disp) | |
| video_imgs.append(np.array(Image.open(im_path))) | |
| disp_video = np.array(video_depths)[:, None] # [:, 0:1, :, :, 0] | |
| img_video = np.array(video_imgs)[..., 0:3] # [:, 0, :, :, 0:3] | |
| print(disp_video.max(), disp_video.min()) | |
| def even_or_odd(num): | |
| if num % 2 == 0: | |
| return num | |
| else: | |
| return num - 1 | |
| # print(disp_video.shape) | |
| # print(img_video.shape) | |
| height = disp_video.shape[-2] | |
| width = disp_video.shape[-1] | |
| height = even_or_odd(height) | |
| width = even_or_odd(width) | |
| disp_video = disp_video[:, :, 0:height, 0:width] | |
| img_video = img_video[:, 0:height, 0:width] | |
| 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"{seq_name}_rgb_left.mp4") | |
| disp_video_path = os.path.join(disp_video_dir, f"{seq_name}_disparity.npz") | |
| imageio.mimsave( | |
| img_video_path, img_video, fps=15, quality=9, macro_block_size=1 | |
| ) | |
| np.savez(disp_video_path, disparity=disp_video) | |
| sample = {} | |
| sample["filepath_left"] = os.path.join( | |
| f"{datatset_name}/{seq_name}_rgb_left.mp4" | |
| ) # img_video_path | |
| sample["filepath_disparity"] = os.path.join( | |
| f"{datatset_name}/{seq_name}_disparity.npz" | |
| ) # disp_video_path | |
| all_samples.append(sample) | |
| # save csv file | |
| 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_bonn( | |
| root="path/to/Bonn-RGBD", | |
| depth_root="path/to/Bonn-RGBD", | |
| saved_rgb_dir="./benchmark/datasets/", | |
| saved_disp_dir="./benchmark/datasets/", | |
| csv_save_path=f"./benchmark/datasets/bonn.csv", | |
| sample_len=-1, | |
| datatset_name="bonn", | |
| start_frame=30, | |
| end_frame=140, | |
| ) | |