Spaces:
Running
on
L4
Running
on
L4
File size: 3,662 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 |
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
from tqdm import tqdm
import matplotlib.pyplot as plt
import shutil
import src.dust3r.datasets.utils.cropping as cropping # noqa
from scipy.spatial.transform import Rotation as R
def get_parser():
import argparse
parser = argparse.ArgumentParser()
parser.add_argument(
"--unreal4k_dir",
default="",
)
parser.add_argument(
"--output_dir",
default="",
)
return parser
def parse_extrinsics(file_path):
"""
Parse the extrinsics file to extract the intrinsics and pose matrices.
Args:
file_path (str): The path to the file containing the extrinsics data.
Returns:
tuple: A tuple containing the intrinsics matrix (3x3) and pose matrix (3x4).
"""
with open(file_path, "r") as file:
lines = file.readlines()
# Parse the intrinsics matrix
intrinsics_data = list(map(float, lines[0].strip().split()))
intrinsics_matrix = np.array(intrinsics_data).reshape(3, 3)
# Parse the pose matrix
cam2world = np.eye(4)
pose_data = list(map(float, lines[1].strip().split()))
pose_matrix = np.array(pose_data).reshape(3, 4)
cam2world[:3] = pose_matrix
cam2world = np.linalg.inv(cam2world)
return intrinsics_matrix, cam2world
def main(rootdir, outdir):
os.makedirs(outdir, exist_ok=True)
envs = [
f for f in sorted(os.listdir(rootdir)) if os.path.isdir(osp.join(rootdir, f))
]
for env in tqdm(envs):
subscenes = ["0", "1"]
for subscene in tqdm(subscenes):
frame_dir = osp.join(rootdir, env)
rgb_dir = osp.join(frame_dir, f"Image{subscene}")
disp_dir = osp.join(frame_dir, f"Disp{subscene}")
ext_dir = osp.join(frame_dir, f"Extrinsics{subscene}")
frame_num = len(os.listdir(rgb_dir))
os.makedirs(osp.join(outdir, env, subscene), exist_ok=True)
for i in tqdm(range(frame_num)):
rgb_path = osp.join(rgb_dir, f"{i:05d}.png")
out_rgb_path = osp.join(outdir, env, subscene, f"{i:05d}_rgb.png")
disp_path = osp.join(disp_dir, f"{i:05d}.npy")
out_depth_path = osp.join(outdir, env, subscene, f"{i:05d}_depth.npy")
out_cam_path = osp.join(outdir, env, subscene, f"{i:05d}.npz")
ext_path0 = osp.join(frame_dir, f"Extrinsics0", f"{i:05d}.txt")
ext_path1 = osp.join(frame_dir, f"Extrinsics1", f"{i:05d}.txt")
K0, c2w0 = parse_extrinsics(ext_path0)
K1, c2w1 = parse_extrinsics(ext_path1)
if subscene == "0":
K = K0
c2w = c2w0
else:
K = K1
c2w = c2w1
img = Image.open(rgb_path).convert("RGB")
disp = np.load(disp_path).astype(np.float32)
baseline = (np.linalg.inv(c2w0) @ c2w1)[0, 3]
depth = baseline * K[0, 0] / disp
image, depthmap, camera_intrinsics = cropping.rescale_image_depthmap(
img, depth, K, output_resolution=(512, 384)
)
image.save(out_rgb_path)
np.save(out_depth_path, depthmap)
np.savez(out_cam_path, intrinsics=camera_intrinsics, cam2world=c2w)
if __name__ == "__main__":
parser = get_parser()
args = parser.parse_args()
main(args.unreal4k_dir, args.output_dir)
|