115 lines
3.3 KiB
Python
115 lines
3.3 KiB
Python
import argparse
|
|
import os
|
|
import re
|
|
from glob import glob
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
|
|
|
|
def raw_camparam_from_xml(path, pose="lookAt"):
|
|
import xml.etree.ElementTree as ET
|
|
|
|
tree = ET.parse(path)
|
|
elm = tree.find("./sensor/transform/" + pose)
|
|
camparam = elm.attrib
|
|
origin = np.fromstring(camparam["origin"], dtype=np.float32, sep=",")
|
|
target = np.fromstring(camparam["target"], dtype=np.float32, sep=",")
|
|
up = np.fromstring(camparam["up"], dtype=np.float32, sep=",")
|
|
height = int(tree.find("./sensor/film/integer[@name='height']").attrib["value"])
|
|
width = int(tree.find("./sensor/film/integer[@name='width']").attrib["value"])
|
|
|
|
camparam = dict()
|
|
camparam["origin"] = origin
|
|
camparam["up"] = up
|
|
camparam["target"] = target
|
|
camparam["height"] = height
|
|
camparam["width"] = width
|
|
return camparam
|
|
|
|
|
|
def get_cam_pos(origin, target, up):
|
|
inward = origin - target
|
|
right = np.cross(up, inward)
|
|
up = np.cross(inward, right)
|
|
rx = np.cross(up, inward)
|
|
ry = np.array(up)
|
|
rz = np.array(inward)
|
|
rx /= np.linalg.norm(rx)
|
|
ry /= np.linalg.norm(ry)
|
|
rz /= np.linalg.norm(rz)
|
|
|
|
rot = np.stack([rx, ry, -rz], axis=0)
|
|
|
|
aff = np.concatenate([np.eye(3), -origin[:, None]], axis=1)
|
|
|
|
ext = np.matmul(rot, aff)
|
|
|
|
result = np.concatenate([ext, np.array([[0, 0, 0, 1]])], axis=0)
|
|
|
|
return result
|
|
|
|
|
|
def convert_cam_params_all_views(datapoint_dir, dataroot, camera_param_dir):
|
|
depths = sorted(glob(os.path.join(datapoint_dir, "*depth.png")))
|
|
cam_ext = [
|
|
"_".join(re.sub(dataroot.strip("/"), camera_param_dir.strip("/"), f).split("_")[:-1]) + ".xml" for f in depths
|
|
]
|
|
|
|
for i, (f, pth) in enumerate(zip(cam_ext, depths)):
|
|
if not os.path.exists(f):
|
|
continue
|
|
params = raw_camparam_from_xml(f)
|
|
origin, target, up, width, height = (
|
|
params["origin"],
|
|
params["target"],
|
|
params["up"],
|
|
params["width"],
|
|
params["height"],
|
|
)
|
|
|
|
ext_matrix = get_cam_pos(origin, target, up)
|
|
|
|
#####
|
|
diag = (0.036**2 + 0.024**2) ** 0.5
|
|
focal_length = 0.05
|
|
res = [480, 480]
|
|
h_relative = res[1] / res[0]
|
|
sensor_width = np.sqrt(diag**2 / (1 + h_relative**2))
|
|
pix_size = sensor_width / res[0]
|
|
|
|
K = np.array(
|
|
[
|
|
[focal_length / pix_size, 0, (sensor_width / pix_size - 1) / 2],
|
|
[0, -focal_length / pix_size, (sensor_width * (res[1] / res[0]) / pix_size - 1) / 2],
|
|
[0, 0, 1],
|
|
]
|
|
)
|
|
|
|
np.savez(pth.split("depth.png")[0] + "cam_params.npz", extr=ext_matrix, intr=K)
|
|
|
|
|
|
def main(opt):
|
|
dataroot_dir = Path(opt.dataroot)
|
|
|
|
leaf_subdirs = []
|
|
|
|
for dirpath, dirnames, filenames in os.walk(dataroot_dir):
|
|
if (not dirnames) and opt.mitsuba_xml_root not in dirpath:
|
|
leaf_subdirs.append(dirpath)
|
|
|
|
for k, dir_ in enumerate(leaf_subdirs):
|
|
print("Processing dir {}/{}: {}".format(k, len(leaf_subdirs), dir_))
|
|
|
|
convert_cam_params_all_views(dir_, opt.dataroot, opt.mitsuba_xml_root)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
args = argparse.ArgumentParser()
|
|
args.add_argument("--dataroot", type=str, default="GenReData/")
|
|
args.add_argument("--mitsuba_xml_root", type=str, default="GenReData/genre-xml_v2")
|
|
|
|
opt = args.parse_args()
|
|
|
|
main(opt)
|