95 lines
3.2 KiB
Python
95 lines
3.2 KiB
Python
|
"""Given the output of the run_colmap.py script, create a usable dataset
|
||
|
|
compatible with the NeRF++ format. Made for Unix-like (posix style path), has
|
||
|
|
not been tested for Windows."""
|
||
|
|
|
||
|
|
import os
|
||
|
|
import json
|
||
|
|
import numpy as np
|
||
|
|
import shutil
|
||
|
|
|
||
|
|
input_path = "./lupo_output_dir_4"
|
||
|
|
output_path = "./lupo_output_data_npp_3"
|
||
|
|
|
||
|
|
# Read the JSON file containing all the data
|
||
|
|
cam_dict_norm_path = os.path.join(input_path, "posed_images/kai_cameras_normalized.json")
|
||
|
|
with open(cam_dict_norm_path) as fd:
|
||
|
|
cam_dict_norm = json.load(fd)
|
||
|
|
|
||
|
|
# Make the train directories
|
||
|
|
train_path = os.path.join(output_path, "train")
|
||
|
|
train_int_path = os.path.join(train_path, "intrinsics")
|
||
|
|
train_pose_path = os.path.join(train_path, "pose")
|
||
|
|
train_rgb_path = os.path.join(train_path, "rgb")
|
||
|
|
|
||
|
|
os.makedirs(train_path, exist_ok=True)
|
||
|
|
os.makedirs(train_int_path, exist_ok=True)
|
||
|
|
os.makedirs(train_pose_path, exist_ok=True)
|
||
|
|
os.makedirs(train_rgb_path, exist_ok=True)
|
||
|
|
|
||
|
|
# Make the train directories
|
||
|
|
test_path = os.path.join(output_path, "test")
|
||
|
|
test_int_path = os.path.join(test_path, "intrinsics")
|
||
|
|
test_pose_path = os.path.join(test_path, "pose")
|
||
|
|
test_rgb_path = os.path.join(test_path, "rgb")
|
||
|
|
|
||
|
|
os.makedirs(test_path, exist_ok=True)
|
||
|
|
os.makedirs(test_int_path, exist_ok=True)
|
||
|
|
os.makedirs(test_pose_path, exist_ok=True)
|
||
|
|
os.makedirs(test_rgb_path, exist_ok=True)
|
||
|
|
|
||
|
|
# Sample images for the test set
|
||
|
|
N = 10
|
||
|
|
image_list = sorted(cam_dict_norm.keys())
|
||
|
|
sampled = np.random.choice(image_list, N, replace=False)
|
||
|
|
|
||
|
|
# Write the files with the corresponding data
|
||
|
|
for img_name in sorted(cam_dict_norm.keys()):
|
||
|
|
# Retrieve the data
|
||
|
|
K = np.array(cam_dict_norm[img_name]['K'])
|
||
|
|
W2C = np.array(cam_dict_norm[img_name]['W2C'])
|
||
|
|
|
||
|
|
img_name_path = os.path.splitext(img_name)[0] # properly espace the extension
|
||
|
|
|
||
|
|
# training set
|
||
|
|
if img_name not in sampled:
|
||
|
|
# Create the paths
|
||
|
|
train_int_img_path = os.path.join(train_int_path, img_name_path + ".txt")
|
||
|
|
train_pose_img_path = os.path.join(train_pose_path, img_name_path + ".txt")
|
||
|
|
train_rgb_img_path = os.path.join(train_rgb_path, img_name)
|
||
|
|
|
||
|
|
# Write intrinsics
|
||
|
|
with open(train_int_img_path, "w") as fd:
|
||
|
|
fd.write(" ".join(map(str, K)))
|
||
|
|
|
||
|
|
# Write poses
|
||
|
|
with open(train_pose_img_path, "w") as fd:
|
||
|
|
fd.write(" ".join(map(str, W2C)))
|
||
|
|
|
||
|
|
# Copy image
|
||
|
|
source_image_path = os.path.join(input_path, "posed_images/images", img_name)
|
||
|
|
shutil.copy(source_image_path, train_rgb_img_path)
|
||
|
|
|
||
|
|
# testing set
|
||
|
|
else:
|
||
|
|
# Create the paths
|
||
|
|
test_int_img_path = os.path.join(test_int_path, img_name_path + ".txt")
|
||
|
|
test_pose_img_path = os.path.join(test_pose_path, img_name_path + ".txt")
|
||
|
|
test_rgb_img_path = os.path.join(test_rgb_path, img_name)
|
||
|
|
|
||
|
|
# Write intrinsics
|
||
|
|
with open(test_int_img_path, "w") as fd:
|
||
|
|
fd.write(" ".join(map(str, K)))
|
||
|
|
|
||
|
|
# Write poses
|
||
|
|
with open(test_pose_img_path, "w") as fd:
|
||
|
|
fd.write(" ".join(map(str, W2C)))
|
||
|
|
|
||
|
|
# Copy image
|
||
|
|
source_image_path = os.path.join(input_path, "posed_images/images", img_name)
|
||
|
|
shutil.copy(source_image_path, test_rgb_img_path)
|
||
|
|
|
||
|
|
|
||
|
|
# Create the validation dataset
|
||
|
|
validation_path = os.path.join(output_path, "validation")
|
||
|
|
os.symlink("./test", validation_path)
|