| ''' |
| ----------------------------------------------------------------------------- |
| Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. |
| |
| NVIDIA CORPORATION and its licensors retain all intellectual property |
| and proprietary rights in and to this software, related documentation |
| and any modifications thereto. Any use, reproduction, disclosure or |
| distribution of this software and related documentation without an express |
| license agreement from NVIDIA CORPORATION is strictly prohibited. |
| ----------------------------------------------------------------------------- |
| ''' |
|
|
| import os |
| import sys |
| from argparse import ArgumentParser |
| from pathlib import Path |
| import yaml |
| from addict import Dict |
| from PIL import Image, ImageFile |
|
|
| dir_path = Path(os.path.dirname(os.path.realpath(__file__))).parents[2] |
| sys.path.append(dir_path.__str__()) |
|
|
| ImageFile.LOAD_TRUNCATED_IMAGES = True |
|
|
|
|
| def generate_config(args): |
| cfg = Dict() |
| cfg._parent_ = "projects/neuralangelo/configs/base.yaml" |
| num_images = len(os.listdir(os.path.join(args.data_dir, "images"))) |
| |
| if args.auto_exposure_wb: |
| cfg.data.num_images = num_images |
| cfg.model.appear_embed.enabled = True |
| cfg.model.appear_embed.dim = 8 |
| else: |
| cfg.model.appear_embed.enabled = False |
| if args.scene_type == "outdoor": |
| cfg.model.object.sdf.mlp.inside_out = False |
| cfg.model.object.sdf.encoding.coarse2fine.init_active_level = 8 |
| elif args.scene_type == "indoor": |
| cfg.model.object.sdf.mlp.inside_out = True |
| cfg.model.object.sdf.encoding.coarse2fine.init_active_level = 8 |
| cfg.model.background.enabled = False |
| cfg.model.render.num_samples.background = 0 |
| elif args.scene_type == "object": |
| cfg.model.object.sdf.mlp.inside_out = False |
| cfg.model.object.sdf.encoding.coarse2fine.init_active_level = 4 |
| else: |
| raise TypeError("Unknown scene type") |
| |
| cfg.data.type = "projects.neuralangelo.data" |
| cfg.data.root = args.data_dir |
| img = Image.open(os.path.join(args.data_dir, "images", os.listdir(os.path.join(args.data_dir, "images"))[0])) |
| w, h = img.size |
| cfg.data.train.image_size = [h, w] |
| short_size = args.val_short_size |
| cfg.data.val.image_size = [short_size, int(w/h*short_size)] if w > h else [int(h/w*short_size), short_size] |
| cfg.data.readjust.center = [0., 0., 0.] |
| cfg.data.readjust.scale = 1. |
| |
| cfg_fname = os.path.join(dir_path, "projects/neuralangelo/configs", f"custom/{args.experiment_name}.yaml") |
| with open(cfg_fname, "w") as file: |
| yaml.safe_dump(cfg.to_dict(), file, default_flow_style=False, indent=4) |
| print("Config generated to file: ", cfg_fname) |
| return |
|
|
|
|
| if __name__ == "__main__": |
| parser = ArgumentParser() |
| parser.add_argument("--experiment_name", type=str, default="recon", help="Name of experiment") |
| parser.add_argument("--data_dir", type=str, default=None, help="Path to data") |
| parser.add_argument("--auto_exposure_wb", action="store_true", |
| help="Video capture with auto-exposure or white-balance") |
| parser.add_argument("--scene_type", type=str, default="outdoor", choices=["outdoor", "indoor", "object"], |
| help="Select scene type. Outdoor for building-scale reconstruction; " |
| "indoor for room-scale reconstruction; object for object-centric scene reconstruction.") |
| parser.add_argument("--val_short_size", type=int, default=300, |
| help="Set the short side of validation images (for saving compute when rendering val images)") |
| args = parser.parse_args() |
| generate_config(args) |
|
|
