Spaces:
Running
on
L4
Running
on
L4
| import argparse | |
| import random | |
| import gzip | |
| import json | |
| import os | |
| import os.path as osp | |
| import torch | |
| import PIL.Image | |
| from PIL import Image | |
| import numpy as np | |
| import cv2 | |
| import multiprocessing | |
| from tqdm import tqdm | |
| import matplotlib.pyplot as plt | |
| import shutil | |
| import path_to_root # noqa | |
| import datasets_preprocess.utils.cropping as cropping # noqa | |
| def get_parser(): | |
| parser = argparse.ArgumentParser() | |
| parser.add_argument("--scannet_dir", default="data/data_scannet") | |
| parser.add_argument("--output_dir", default="data/dust3r_data/processed_scannet") | |
| return parser | |
| def process_scene(args): | |
| rootdir, outdir, split, scene = args | |
| frame_dir = osp.join(rootdir, split, scene) | |
| rgb_dir = osp.join(frame_dir, "color") | |
| depth_dir = osp.join(frame_dir, "depth") | |
| pose_dir = osp.join(frame_dir, "pose") | |
| depth_intrinsic = np.loadtxt( | |
| osp.join(frame_dir, "intrinsic", "intrinsic_depth.txt") | |
| )[:3, :3].astype(np.float32) | |
| color_intrinsic = np.loadtxt( | |
| osp.join(frame_dir, "intrinsic", "intrinsic_color.txt") | |
| )[:3, :3].astype(np.float32) | |
| if not np.isfinite(depth_intrinsic).all() or not np.isfinite(color_intrinsic).all(): | |
| return | |
| os.makedirs(osp.join(outdir, split, scene), exist_ok=True) | |
| frame_num = len(os.listdir(rgb_dir)) | |
| assert frame_num == len(os.listdir(depth_dir)) == len(os.listdir(pose_dir)) | |
| out_rgb_dir = osp.join(outdir, split, scene, "color") | |
| out_depth_dir = osp.join(outdir, split, scene, "depth") | |
| out_cam_dir = osp.join(outdir, split, scene, "cam") | |
| os.makedirs(out_rgb_dir, exist_ok=True) | |
| os.makedirs(out_depth_dir, exist_ok=True) | |
| os.makedirs(out_cam_dir, exist_ok=True) | |
| for i in tqdm(range(frame_num)): | |
| rgb_path = osp.join(rgb_dir, f"{i}.jpg") | |
| depth_path = osp.join(depth_dir, f"{i}.png") | |
| pose_path = osp.join(pose_dir, f"{i}.txt") | |
| rgb = Image.open(rgb_path) | |
| depth = cv2.imread(depth_path, cv2.IMREAD_UNCHANGED) | |
| rgb = rgb.resize(depth.shape[::-1], resample=Image.Resampling.LANCZOS) | |
| pose = np.loadtxt(pose_path).reshape(4, 4).astype(np.float32) | |
| if not np.isfinite(pose).all(): | |
| continue | |
| out_rgb_path = osp.join(out_rgb_dir, f"{i:05d}.jpg") | |
| out_depth_path = osp.join(out_depth_dir, f"{i:05d}.png") | |
| out_cam_path = osp.join(out_cam_dir, f"{i:05d}.npz") | |
| np.savez(out_cam_path, intrinsics=depth_intrinsic, pose=pose) | |
| rgb.save(out_rgb_path) | |
| cv2.imwrite(out_depth_path, depth) | |
| def main(rootdir, outdir): | |
| os.makedirs(outdir, exist_ok=True) | |
| splits = ["scans_test", "scans_train"] | |
| pool = multiprocessing.Pool(processes=multiprocessing.cpu_count()) | |
| for split in splits: | |
| scenes = [ | |
| f | |
| for f in os.listdir(os.path.join(rootdir, split)) | |
| if os.path.isdir(osp.join(rootdir, split, f)) | |
| ] | |
| pool.map(process_scene, [(rootdir, outdir, split, scene) for scene in scenes]) | |
| pool.close() | |
| pool.join() | |
| if __name__ == "__main__": | |
| parser = get_parser() | |
| args = parser.parse_args() | |
| main(args.scannet_dir, args.output_dir) | |