Spaces:
Runtime error
Runtime error
File size: 4,324 Bytes
2df809d |
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 111 112 113 114 115 116 117 118 119 120 121 122 123 124 |
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)
|