| import os |
| import json |
| import argparse |
|
|
|
|
| def get_splits_paths(dataset_path): |
| |
| im_path = os.path.join(dataset_path, "images") |
| pc_path = dataset_path |
| splits = ["train", "val"] |
| split_path = os.path.join(dataset_path, "splits") |
| os.makedirs(split_path, exist_ok=True) |
| for split in splits: |
| im_split_path = os.path.join(im_path, split) |
| pc_split_path = os.path.join(pc_path, split).replace( |
| dataset_path, "data/structured3d" |
| ) |
| split_names = [f.name for f in os.scandir(im_split_path) if f.is_dir()] |
| split_dict = {} |
| for name in split_names: |
| cam_split_name_path = os.path.join(im_split_path, name) |
| cam_names = [f.name for f in os.scandir(cam_split_name_path) if f.is_dir()] |
| for cam_name in cam_names: |
| im_split_name_path = os.path.join( |
| cam_split_name_path, cam_name, "color", "prsp" |
| ) |
| co_split_name_path = os.path.join( |
| cam_split_name_path, |
| cam_name, |
| "correspondence", |
| "prsp_correspondence", |
| ) |
| png_files = [ |
| f for f in os.listdir(im_split_name_path) if f.endswith(".png") |
| ] |
| png_files = sorted(png_files, key=lambda x: int(x.split(".")[0])) |
| |
| png_file_paths = [ |
| os.path.join(im_split_name_path, f).replace( |
| dataset_path, "data/structured3d" |
| ) |
| for f in png_files |
| ] |
| co_file_paths = [ |
| os.path.join(co_split_name_path, f.replace(".png", ".npy")).replace( |
| dataset_path, "data/structured3d" |
| ) |
| for f in png_files |
| ] |
| split_dict[f"{name}_{cam_name}"] = {} |
| split_dict[f"{name}_{cam_name}"]["pointclouds"] = os.path.join( |
| pc_split_path, name, cam_name |
| ) |
| split_dict[f"{name}_{cam_name}"]["images"] = png_file_paths |
| split_dict[f"{name}_{cam_name}"]["correspondences"] = co_file_paths |
| with open(os.path.join(split_path, f"{split}.json"), "w") as f: |
| json.dump(split_dict, f, indent=4) |
|
|
|
|
| |
| if __name__ == "__main__": |
| parser = argparse.ArgumentParser() |
| parser.add_argument( |
| "--dataset_root", |
| required=True, |
| help="Path to the Structured3D dataset containing scene folders", |
| ) |
| config = parser.parse_args() |
| get_splits_paths(config.dataset_root) |
|
|
