From 5a76b955cdd833fd59b90edff33875fa05894847 Mon Sep 17 00:00:00 2001 From: Ethan Weber Date: Mon, 3 Oct 2022 12:02:10 -0700 Subject: [PATCH] adding --data alias (#630) * adding --data alias * fix black --- nerfstudio/configs/base_config.py | 5 ++++- nerfstudio/configs/model_configs.py | 5 +++-- nerfstudio/data/dataparsers/base_dataparser.py | 2 ++ nerfstudio/data/dataparsers/blender_dataparser.py | 10 +++++----- nerfstudio/data/dataparsers/friends_dataparser.py | 10 +++++----- nerfstudio/data/dataparsers/instant_ngp_dataparser.py | 8 ++++---- nerfstudio/data/dataparsers/mipnerf_dataparser.py | 8 ++++---- nerfstudio/data/dataparsers/nerfstudio_dataparser.py | 10 +++++----- nerfstudio/data/dataparsers/record3d_dataparser.py | 6 +++--- nerfstudio/viewer/server/subprocess.py | 2 +- scripts/train.py | 3 +++ tests/test_run_train.py | 2 +- 12 files changed, 40 insertions(+), 31 deletions(-) diff --git a/nerfstudio/configs/base_config.py b/nerfstudio/configs/base_config.py index 89b8aedd8..2002f3c9d 100644 --- a/nerfstudio/configs/base_config.py +++ b/nerfstudio/configs/base_config.py @@ -206,6 +206,9 @@ class Config(PrintableConfig): ) """optionally specify a pre-defined config to load from""" vis: Literal["viewer", "wandb", "tensorboard"] = "wandb" + """Which visualizer to use.""" + data: Optional[Path] = None + """Alias for --pipeline.datamanager.dataparser.data""" def is_viewer_enabled(self) -> bool: """Checks if a viewer is enabled.""" @@ -226,7 +229,7 @@ class Config(PrintableConfig): def set_experiment_name(self) -> None: """Dynamically set the experiment name""" if self.experiment_name is None: - self.experiment_name = str(self.pipeline.datamanager.dataparser.data_directory).replace("/", "-") + self.experiment_name = str(self.pipeline.datamanager.dataparser.data).replace("/", "-") def get_base_dir(self) -> Path: """Retrieve the base directory to set relative paths""" diff --git a/nerfstudio/configs/model_configs.py b/nerfstudio/configs/model_configs.py index af96a3fc1..b97b69c2a 100644 --- a/nerfstudio/configs/model_configs.py +++ b/nerfstudio/configs/model_configs.py @@ -27,6 +27,7 @@ from nerfstudio.data.datamanagers import VanillaDataManagerConfig from nerfstudio.data.dataparsers.blender_dataparser import BlenderDataParserConfig from nerfstudio.data.dataparsers.friends_dataparser import FriendsDataParserConfig from nerfstudio.data.dataparsers.mipnerf_dataparser import MipNerf360DataParserConfig +from nerfstudio.data.dataparsers.nerfstudio_dataparser import NerfstudioDataParserConfig from nerfstudio.engine.optimizers import AdamOptimizerConfig, RAdamOptimizerConfig from nerfstudio.engine.schedulers import SchedulerConfig from nerfstudio.models.base_model import VanillaModelConfig @@ -59,7 +60,7 @@ model_configs["nerfacto"] = Config( trainer=TrainerConfig(steps_per_eval_batch=500, steps_per_save=2000, mixed_precision=True), pipeline=VanillaPipelineConfig( datamanager=VanillaDataManagerConfig( - dataparser=BlenderDataParserConfig(), train_num_rays_per_batch=4096, eval_num_rays_per_batch=8192 + dataparser=NerfstudioDataParserConfig(), train_num_rays_per_batch=4096, eval_num_rays_per_batch=8192 ), model=NerfactoModelConfig(eval_num_rays_per_chunk=1 << 16), ), @@ -81,7 +82,7 @@ model_configs["instant-ngp"] = Config( method_name="instant-ngp", trainer=TrainerConfig(steps_per_eval_batch=500, steps_per_save=2000, mixed_precision=True), pipeline=DynamicBatchPipelineConfig( - datamanager=VanillaDataManagerConfig(dataparser=BlenderDataParserConfig(), train_num_rays_per_batch=8192), + datamanager=VanillaDataManagerConfig(dataparser=NerfstudioDataParserConfig(), train_num_rays_per_batch=8192), model=InstantNGPModelConfig(eval_num_rays_per_chunk=8192), ), optimizers={ diff --git a/nerfstudio/data/dataparsers/base_dataparser.py b/nerfstudio/data/dataparsers/base_dataparser.py index ce083d5f7..102342607 100644 --- a/nerfstudio/data/dataparsers/base_dataparser.py +++ b/nerfstudio/data/dataparsers/base_dataparser.py @@ -80,6 +80,8 @@ class DataParserConfig(cfg.InstantiateConfig): _target: Type = field(default_factory=lambda: DataParser) """_target: target class to instantiate""" + data: Path = Path() + """Directory specifying location of data.""" @dataclass diff --git a/nerfstudio/data/dataparsers/blender_dataparser.py b/nerfstudio/data/dataparsers/blender_dataparser.py index 7f6a62a13..5a653537e 100644 --- a/nerfstudio/data/dataparsers/blender_dataparser.py +++ b/nerfstudio/data/dataparsers/blender_dataparser.py @@ -40,8 +40,8 @@ class BlenderDataParserConfig(DataParserConfig): _target: Type = field(default_factory=lambda: Blender) """target class to instantiate""" - data_directory: Path = Path("data/blender/lego") - """directory specifying location of data""" + data: Path = Path("data/blender/lego") + """Directory specifying location of data.""" scale_factor: float = 1.0 """How much to scale the camera origins by.""" alpha_color: str = "white" @@ -58,7 +58,7 @@ class Blender(DataParser): def __init__(self, config: BlenderDataParserConfig): super().__init__(config=config) - self.data_directory: Path = config.data_directory + self.data: Path = config.data self.scale_factor: float = config.scale_factor self.alpha_color = config.alpha_color @@ -68,11 +68,11 @@ class Blender(DataParser): else: alpha_color_tensor = None - meta = load_from_json(self.data_directory / f"transforms_{split}.json") + meta = load_from_json(self.data / f"transforms_{split}.json") image_filenames = [] poses = [] for frame in meta["frames"]: - fname = self.data_directory / Path(frame["file_path"].replace("./", "") + ".png") + fname = self.data / Path(frame["file_path"].replace("./", "") + ".png") image_filenames.append(fname) poses.append(np.array(frame["transform_matrix"])) poses = np.array(poses).astype(np.float32) diff --git a/nerfstudio/data/dataparsers/friends_dataparser.py b/nerfstudio/data/dataparsers/friends_dataparser.py index 1926b4afb..3ec052d5d 100644 --- a/nerfstudio/data/dataparsers/friends_dataparser.py +++ b/nerfstudio/data/dataparsers/friends_dataparser.py @@ -63,8 +63,8 @@ class FriendsDataParserConfig(DataParserConfig): _target: Type = field(default_factory=lambda: Friends) """target class to instantiate""" - data_directory: Path = Path("data/friends/TBBT-big_living_room") - """directory specifying location of data""" + data: Path = Path("data/friends/TBBT-big_living_room") + """Directory specifying location of data.""" include_semantics: bool = True """whether or not to include loading of semantics data""" downscale_factor: int = 8 @@ -83,7 +83,7 @@ class Friends(DataParser): def _generate_dataparser_outputs(self, split="train"): # pylint: disable=unused-argument,too-many-statements - cameras_json = load_from_json(self.config.data_directory / "cameras.json") + cameras_json = load_from_json(self.config.data / "cameras.json") frames = cameras_json["frames"] bbox = torch.tensor(cameras_json["bbox"]) @@ -99,7 +99,7 @@ class Friends(DataParser): camera_to_worlds = [] for frame in frames: # unpack data - image_filename = self.config.data_directory / images_folder / frame["image_name"] + image_filename = self.config.data / images_folder / frame["image_name"] intrinsics = torch.tensor(frame["intrinsics"]) camtoworld = torch.tensor(frame["camtoworld"])[:3] # append data @@ -155,7 +155,7 @@ class Friends(DataParser): ) for image_filename in image_filenames ] - panoptic_classes = load_from_json(self.config.data_directory / "panoptic_classes.json") + panoptic_classes = load_from_json(self.config.data / "panoptic_classes.json") stuff_classes = panoptic_classes["stuff"] stuff_colors = torch.tensor(panoptic_classes["stuff_colors"], dtype=torch.float32) / 255.0 thing_classes = panoptic_classes["thing"] diff --git a/nerfstudio/data/dataparsers/instant_ngp_dataparser.py b/nerfstudio/data/dataparsers/instant_ngp_dataparser.py index a5578aa62..64244fdab 100644 --- a/nerfstudio/data/dataparsers/instant_ngp_dataparser.py +++ b/nerfstudio/data/dataparsers/instant_ngp_dataparser.py @@ -41,8 +41,8 @@ class InstantNGPDataParserConfig(DataParserConfig): _target: Type = field(default_factory=lambda: InstantNGP) """target class to instantiate""" - data_directory: Path = Path("data/ours/posterv2") - """directory specifying location of data""" + data: Path = Path("data/ours/posterv2") + """Directory specifying location of data.""" scale_factor: float = 1.0 """How much to scale the camera origins by.""" scene_scale: float = 0.33 @@ -57,12 +57,12 @@ class InstantNGP(DataParser): def _generate_dataparser_outputs(self, split="train"): - meta = load_from_json(self.config.data_directory / "transforms.json") + meta = load_from_json(self.config.data / "transforms.json") image_filenames = [] poses = [] num_skipped_image_filenames = 0 for frame in meta["frames"]: - fname = self.config.data_directory / Path(frame["file_path"]) + fname = self.config.data / Path(frame["file_path"]) if not fname: num_skipped_image_filenames += 1 else: diff --git a/nerfstudio/data/dataparsers/mipnerf_dataparser.py b/nerfstudio/data/dataparsers/mipnerf_dataparser.py index 71d69695b..3f320759b 100644 --- a/nerfstudio/data/dataparsers/mipnerf_dataparser.py +++ b/nerfstudio/data/dataparsers/mipnerf_dataparser.py @@ -38,8 +38,8 @@ class MipNerf360DataParserConfig(DataParserConfig): _target: Type = field(default_factory=lambda: Mipnerf360) """target class to instantiate""" - data_directory: Path = Path("data/mipnerf_360/garden") - """directory specifying location of data""" + data: Path = Path("data/mipnerf_360/garden") + """Directory specifying location of data.""" downscale_factor: int = 1 """How much to downscale images.""" val_skip: int = 8 @@ -82,7 +82,7 @@ class Mipnerf360(DataParser): image_dir = "images" if self.config.downscale_factor > 1: image_dir += f"_{self.config.downscale_factor}" - image_dir = self.config.data_directory / image_dir + image_dir = self.config.data / image_dir if not image_dir.exists(): raise ValueError(f"Image directory {image_dir} doesn't exist") @@ -96,7 +96,7 @@ class Mipnerf360(DataParser): image_filenames = sorted(image_filenames) num_images = len(image_filenames) - poses_data = np.load(self.config.data_directory / "poses_bounds.npy") + poses_data = np.load(self.config.data / "poses_bounds.npy") poses = poses_data[:, :-2].reshape([-1, 3, 5]).astype(np.float32) bounds = poses_data[:, -2:].transpose([1, 0]) diff --git a/nerfstudio/data/dataparsers/nerfstudio_dataparser.py b/nerfstudio/data/dataparsers/nerfstudio_dataparser.py index b78c5dc9d..5a8e1fbc9 100644 --- a/nerfstudio/data/dataparsers/nerfstudio_dataparser.py +++ b/nerfstudio/data/dataparsers/nerfstudio_dataparser.py @@ -41,8 +41,8 @@ class NerfstudioDataParserConfig(DataParserConfig): _target: Type = field(default_factory=lambda: Nerfstudio) """target class to instantiate""" - data_directory: Path = Path("data/nerfstudio/poster") - """directory specifying location of data""" + data: Path = Path("data/nerfstudio/poster") + """Directory specifying location of data.""" scale_factor: float = 1.0 """How much to scale the camera origins by.""" downscale_factor: int = 1 @@ -65,7 +65,7 @@ class Nerfstudio(DataParser): def _generate_dataparser_outputs(self, split="train"): # pylint: disable=too-many-statements - meta = load_from_json(self.config.data_directory / "transforms.json") + meta = load_from_json(self.config.data / "transforms.json") image_filenames = [] poses = [] num_skipped_image_filenames = 0 @@ -75,9 +75,9 @@ class Nerfstudio(DataParser): else: filepath = Path(frame["file_path"]) if self.config.downscale_factor > 1: - fname = self.config.data_directory / f"images_{self.config.downscale_factor}" / filepath.name + fname = self.config.data / f"images_{self.config.downscale_factor}" / filepath.name else: - fname = self.config.data_directory / filepath + fname = self.config.data / filepath if not fname: num_skipped_image_filenames += 1 else: diff --git a/nerfstudio/data/dataparsers/record3d_dataparser.py b/nerfstudio/data/dataparsers/record3d_dataparser.py index 4edc5e4c3..484d1ddae 100644 --- a/nerfstudio/data/dataparsers/record3d_dataparser.py +++ b/nerfstudio/data/dataparsers/record3d_dataparser.py @@ -41,7 +41,7 @@ class Record3DDataParserConfig(DataParserConfig): _target: Type = field(default_factory=lambda: Record3D) """target class to instantiate""" - data_directory: Path = Path("data/record3d/garden") + data: Path = Path("data/record3d/garden") """Location of data""" val_skip: int = 8 """1/val_skip images to use for validation.""" @@ -61,7 +61,7 @@ class Record3D(DataParser): config: Record3DDataParserConfig def _generate_dataparser_outputs(self, split: str = "train") -> DataparserOutputs: - image_dir = self.config.data_directory / "rgb" + image_dir = self.config.data / "rgb" if not image_dir.exists(): raise ValueError(f"Image directory {image_dir} doesn't exist") @@ -75,7 +75,7 @@ class Record3D(DataParser): image_filenames = np.array(image_filenames) num_images = len(image_filenames) - metadata_path = self.config.data_directory / "metadata.json" + metadata_path = self.config.data / "metadata.json" metadata_dict = load_from_json(metadata_path) poses_data = np.array(metadata_dict["poses"]) diff --git a/nerfstudio/viewer/server/subprocess.py b/nerfstudio/viewer/server/subprocess.py index 6d902ba28..5e5db53be 100644 --- a/nerfstudio/viewer/server/subprocess.py +++ b/nerfstudio/viewer/server/subprocess.py @@ -54,7 +54,7 @@ def run_viewer_bridge_server_as_subprocess( sock.bind(("", 0)) zmq_port = sock.getsockname()[1] string = f"Using ZMQ port: {zmq_port}" - CONSOLE.print(f"[bold red]{string}") + CONSOLE.print(f"[bold yellow]{string}") args.append("--zmq-port") args.append(str(zmq_port)) diff --git a/scripts/train.py b/scripts/train.py index fbabee9c3..0f6fbb235 100755 --- a/scripts/train.py +++ b/scripts/train.py @@ -187,6 +187,9 @@ def main(config: cfg.Config) -> None: """Main function.""" config.set_timestamp() + if config.data: + logging.info("Using --data alias for --data.pipeline.datamanager.dataparser.data") + config.pipeline.datamanager.dataparser.data = config.data if config.trainer.load_config: logging.info(f"Loading pre-set config from: {config.trainer.load_config}") diff --git a/tests/test_run_train.py b/tests/test_run_train.py index 2c694e006..9b5e19dd6 100644 --- a/tests/test_run_train.py +++ b/tests/test_run_train.py @@ -22,7 +22,7 @@ def set_reduced_config(config: Config): config.machine.num_gpus = 0 config.trainer.max_num_iterations = 2 # reduce dataset factors; set dataset to test - config.pipeline.datamanager.dataparser = BlenderDataParserConfig(data_directory=Path("tests/data/lego_test")) + config.pipeline.datamanager.dataparser = BlenderDataParserConfig(data=Path("tests/data/lego_test")) config.pipeline.datamanager.train_num_images_to_sample_from = 1 config.pipeline.datamanager.train_num_rays_per_batch = 4