Spaces:
Running
on
L4
Running
on
L4
| import os | |
| import os.path as osp | |
| from PIL import Image | |
| import numpy as np | |
| from tqdm import tqdm | |
| from read_write_model import run | |
| def get_parser(): | |
| import argparse | |
| parser = argparse.ArgumentParser() | |
| parser.add_argument("--mapfree_dir", default="") # TODO | |
| parser.add_argument("--output_dir", default="test_preprocess") # TODO | |
| return parser | |
| def main(rootdir, outdir): | |
| os.makedirs(outdir, exist_ok=True) | |
| envs = [f for f in os.listdir(rootdir) if os.path.isdir(osp.join(rootdir, f))] | |
| for env in tqdm(envs): | |
| subseqs = [ | |
| f | |
| for f in os.listdir(osp.join(rootdir, env)) | |
| if os.path.isdir(osp.join(rootdir, env, f)) | |
| ] | |
| for subseq in subseqs: | |
| sparse_dir = osp.join(rootdir, env, subseq, "sparse") | |
| images_dir = osp.join(rootdir, env, subseq, "images") | |
| run(sparse_dir, sparse_dir) | |
| intrins_file = sparse_dir + "/cameras.txt" | |
| poses_file = sparse_dir + "/images.txt" | |
| cam_params = {} | |
| with open(intrins_file, "r") as f: | |
| for line in f: | |
| if line.startswith("#"): | |
| continue | |
| parts = line.strip().split() | |
| if len(parts) == 0: | |
| continue | |
| cam_id = int(parts[0]) | |
| fx = float(parts[4]) | |
| fy = float(parts[5]) | |
| cx = float(parts[6]) | |
| cy = float(parts[7]) | |
| cam_params[cam_id] = np.array([[fx, 0, cx], [0, fy, cy], [0, 0, 1]]) | |
| poses = [] | |
| images = [] | |
| intrinsics = [] | |
| with open(poses_file, "r") as f: | |
| for i, line in enumerate(f): | |
| if line.startswith("#"): | |
| continue | |
| parts = line.strip().split() | |
| if len(parts) == 0: | |
| continue | |
| if "." in parts[0]: | |
| continue | |
| img_name = parts[-1] | |
| w, x, y, z = map(float, parts[1:5]) | |
| R = np.array( | |
| [ | |
| [ | |
| 1 - 2 * y * y - 2 * z * z, | |
| 2 * x * y - 2 * z * w, | |
| 2 * x * z + 2 * y * w, | |
| ], | |
| [ | |
| 2 * x * y + 2 * z * w, | |
| 1 - 2 * x * x - 2 * z * z, | |
| 2 * y * z - 2 * x * w, | |
| ], | |
| [ | |
| 2 * x * z - 2 * y * w, | |
| 2 * y * z + 2 * x * w, | |
| 1 - 2 * x * x - 2 * y * y, | |
| ], | |
| ] | |
| ) | |
| tx, ty, tz = map(float, parts[5:8]) | |
| cam_id = int(parts[-2]) | |
| pose = np.eye(4) | |
| pose[:3, :3] = R | |
| pose[:3, 3] = [tx, ty, tz] | |
| poses.append(np.linalg.inv(pose)) | |
| images.append(img_name) | |
| intrinsics.append(cam_params[cam_id]) | |
| os.makedirs(osp.join(outdir, env, subseq), exist_ok=True) | |
| os.makedirs(osp.join(outdir, env, subseq, "rgb"), exist_ok=True) | |
| os.makedirs(osp.join(outdir, env, subseq, "cam"), exist_ok=True) | |
| for i, img_name in enumerate(tqdm(images)): | |
| img_path = os.path.join(images_dir, img_name) | |
| rgb = Image.open(img_path) | |
| intrinsic = intrinsics[i] | |
| pose = poses[i] | |
| # save all | |
| basename = img_name.split("/")[-1] | |
| rgb.save(osp.join(outdir, env, subseq, "rgb", basename)) | |
| np.savez( | |
| osp.join( | |
| outdir, env, subseq, "cam", basename.replace(".jpg", ".npz") | |
| ), | |
| intrinsic=intrinsic, | |
| pose=pose, | |
| ) | |
| if __name__ == "__main__": | |
| parser = get_parser() | |
| args = parser.parse_args() | |
| main(args.mapfree_dir, args.output_dir) | |