From 00c3ae73b2d0dcf465c6087cac3cc83fdb6c4760 Mon Sep 17 00:00:00 2001 From: Matthew Tancik Date: Sun, 18 Sep 2022 16:27:41 -0700 Subject: [PATCH] Move dataset parser configs to files (#476) * Move parser configs to files * fixed circular import errors Co-authored-by: Ethan Weber --- nerfactory/configs/base.py | 122 ++---------------- .../datamanagers/dataparsers/__init__.py | 13 ++ nerfactory/datamanagers/dataparsers/base.py | 2 + .../dataparsers/blender_parser.py | 21 ++- .../dataparsers/friends_parser.py | 25 +++- .../dataparsers/instant_ngp_parser.py | 19 ++- .../dataparsers/mipnerf_parser.py | 24 +++- .../dataparsers/nerfactory_parser.py | 31 ++++- .../dataparsers/record3d_parser.py | 25 +++- 9 files changed, 159 insertions(+), 123 deletions(-) create mode 100644 nerfactory/datamanagers/dataparsers/__init__.py diff --git a/nerfactory/configs/base.py b/nerfactory/configs/base.py index 439bfbaf9..aaaeb8dd0 100644 --- a/nerfactory/configs/base.py +++ b/nerfactory/configs/base.py @@ -29,12 +29,6 @@ from nerfactory.configs.utils import to_immutable_dict # data instances from nerfactory.datamanagers.base import VanillaDataManager from nerfactory.datamanagers.dataparsers.base import DataParser -from nerfactory.datamanagers.dataparsers.blender_parser import Blender -from nerfactory.datamanagers.dataparsers.friends_parser import Friends -from nerfactory.datamanagers.dataparsers.instant_ngp_parser import InstantNGP -from nerfactory.datamanagers.dataparsers.mipnerf_parser import Mipnerf360 -from nerfactory.datamanagers.dataparsers.nerfactory_parser import Nerfactory -from nerfactory.datamanagers.dataparsers.record3d_parser import Record3D # model instances from nerfactory.models.base import Model @@ -201,109 +195,19 @@ class DataParserConfig(InstantiateConfig): """_target: target class to instantiate""" -@dataclass -class NerfactoryDataParserConfig(DataParserConfig): - """Nerfactory dataset config""" - - _target: Type = Nerfactory - """target class to instantiate""" - data_directory: Path = Path("data/ours/posters_v3") - """directory specifying location of data""" - scale_factor: float = 1.0 - """How much to scale the camera origins by.""" - downscale_factor: int = 1 - """How much to downscale images. Defaults to 1.""" - scene_scale: float = 4.0 - """How much to scale the scene. Defaults to 4.0""" - orientation_method: Literal["pca", "up"] = "up" - """The method to use for orientation. Either "pca" or "up".""" - train_split_percentage: float = 0.9 - """ - The percent of images to use for training. - The remaining images are for eval. Defaults to 0.9. - """ - - -@dataclass -class BlenderDataParserConfig(DataParserConfig): - """Blender dataset config""" - - _target: Type = Blender - """target class to instantiate""" - data_directory: 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" - """alpha color of background""" - - -@dataclass -class FriendsDataParserConfig(DataParserConfig): - """Friends dataset config""" - - _target: Type = Friends - """target class to instantiate""" - data_directory: 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 - scene_scale: float = 4.0 - """ - Sets the bounding cube to have edge length of this size. - The longest dimension of the Friends axis-aligned bbox will be scaled to this value. - """ - - -@dataclass -class MipNerf360DataParserConfig(DataParserConfig): - """Mipnerf 360 dataset config""" - - _target: Type = Mipnerf360 - """target class to instantiate""" - data_directory: Path = Path("data/mipnerf_360/garden") - """directory specifying location of data""" - downscale_factor: int = 1 - """How much to downscale images. Defaults to 1.""" - val_skip: int = 8 - """1/val_skip images to use for validation. Defaults to 8.""" - auto_scale: bool = True - """Scale based on pose bounds. Defaults to True.""" - aabb_scale: float = 4 - """Scene scale, Defaults to 1.0.""" - - -@dataclass -class InstantNGPDataParserConfig(DataParserConfig): - """Instant-NGP dataset config""" - - _target: Type = InstantNGP - """target class to instantiate""" - data_directory: 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 - """How much to scale the scene. Defaults to 0.33""" - - -@dataclass -class Record3DDataParserConfig(DataParserConfig): - """Record3D dataset config""" - - _target: Type = Record3D - """target class to instantiate""" - data_directory: Path = Path("data/record3d/garden") - """Location of data""" - val_skip: int = 8 - """1/val_skip images to use for validation. Defaults to 8.""" - aabb_scale: float = 4.0 - """Scene scale, Defaults to 4.0.""" - max_dataset_size: int = 150 - """Max number of images to train on. If the dataset has - more, images will be sampled approximately evenly. Defaults to 150.""" - +# pylint: disable=wrong-import-position +from nerfactory.datamanagers.dataparsers.blender_parser import BlenderDataParserConfig +from nerfactory.datamanagers.dataparsers.friends_parser import FriendsDataParserConfig +from nerfactory.datamanagers.dataparsers.instant_ngp_parser import ( + InstantNGPDataParserConfig, +) +from nerfactory.datamanagers.dataparsers.mipnerf_parser import ( + MipNerf360DataParserConfig, +) +from nerfactory.datamanagers.dataparsers.nerfactory_parser import ( + NerfactoryDataParserConfig, +) +from nerfactory.datamanagers.dataparsers.record3d_parser import Record3DDataParserConfig AnnotatedDataParserUnion = dcargs.extras.subcommand_type_from_defaults( { diff --git a/nerfactory/datamanagers/dataparsers/__init__.py b/nerfactory/datamanagers/dataparsers/__init__.py new file mode 100644 index 000000000..02655775e --- /dev/null +++ b/nerfactory/datamanagers/dataparsers/__init__.py @@ -0,0 +1,13 @@ +# Copyright 2022 The Plenoptix Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/nerfactory/datamanagers/dataparsers/base.py b/nerfactory/datamanagers/dataparsers/base.py index d8cb43a10..8039ff82a 100644 --- a/nerfactory/datamanagers/dataparsers/base.py +++ b/nerfactory/datamanagers/dataparsers/base.py @@ -31,6 +31,8 @@ class DataParser: config: datasetparser config containing all information needed to instantiate dataset """ + config: cfg.DataParserConfig + def __init__(self, config: cfg.DataParserConfig): super().__init__() self.config = config diff --git a/nerfactory/datamanagers/dataparsers/blender_parser.py b/nerfactory/datamanagers/dataparsers/blender_parser.py index 970c6a1aa..034bd7018 100644 --- a/nerfactory/datamanagers/dataparsers/blender_parser.py +++ b/nerfactory/datamanagers/dataparsers/blender_parser.py @@ -15,8 +15,9 @@ """Data parser for blender dataset""" from __future__ import annotations -from dataclasses import dataclass +from dataclasses import dataclass, field from pathlib import Path +from typing import Type import imageio import numpy as np @@ -30,13 +31,29 @@ from nerfactory.utils.colors import get_color from nerfactory.utils.io import get_absolute_path, load_from_json +@dataclass +class BlenderDataParserConfig(cfg.DataParserConfig): + """Blender dataset parser config""" + + _target: Type = field(default_factory=lambda: Blender) + """target class to instantiate""" + data_directory: 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" + """alpha color of background""" + + @dataclass class Blender(DataParser): """Blender Dataset Some of this code comes from https://github.com/yenchenlin/nerf-pytorch/blob/master/load_blender.py#L37. """ - def __init__(self, config: cfg.BlenderDataParserConfig): + config: cfg.BlenderDataParserConfig + + def __init__(self, config: BlenderDataParserConfig): super().__init__(config=config) self.data_directory: Path = config.data_directory self.scale_factor: float = config.scale_factor diff --git a/nerfactory/datamanagers/dataparsers/friends_parser.py b/nerfactory/datamanagers/dataparsers/friends_parser.py index bbb41187f..197764ed6 100644 --- a/nerfactory/datamanagers/dataparsers/friends_parser.py +++ b/nerfactory/datamanagers/dataparsers/friends_parser.py @@ -15,8 +15,9 @@ """Data parser for friends dataset""" from __future__ import annotations -from dataclasses import dataclass +from dataclasses import dataclass, field from pathlib import Path +from typing import Type import numpy as np import torch @@ -52,13 +53,31 @@ def get_semantics_and_masks(image_idx: int, semantics: Semantics): return {"mask": mask, "semantics": stuff_semantics} +@dataclass +class FriendsDataParserConfig(cfg.DataParserConfig): + """Friends dataset parser config""" + + _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""" + include_semantics: bool = True + """whether or not to include loading of semantics data""" + downscale_factor: int = 8 + scene_scale: float = 4.0 + """ + Sets the bounding cube to have edge length of this size. + The longest dimension of the Friends axis-aligned bbox will be scaled to this value. + """ + + @dataclass class Friends(DataParser): """Friends Dataset""" - config: cfg.FriendsDataParserConfig + config: FriendsDataParserConfig - def _generate_dataset_inputs(self, split="train"): # pylint: disable=unused-argument + def _generate_dataset_inputs(self, split="train"): # pylint: disable=unused-argument,too-many-statements abs_dir = get_absolute_path(self.config.data_directory) diff --git a/nerfactory/datamanagers/dataparsers/instant_ngp_parser.py b/nerfactory/datamanagers/dataparsers/instant_ngp_parser.py index 952b83948..b0e62c3a8 100644 --- a/nerfactory/datamanagers/dataparsers/instant_ngp_parser.py +++ b/nerfactory/datamanagers/dataparsers/instant_ngp_parser.py @@ -17,8 +17,9 @@ from __future__ import annotations import logging -from dataclasses import dataclass +from dataclasses import dataclass, field from pathlib import Path +from typing import Type import numpy as np import torch @@ -31,11 +32,25 @@ from nerfactory.datamanagers.structs import DatasetInputs, SceneBounds from nerfactory.utils.io import get_absolute_path, load_from_json +@dataclass +class InstantNGPDataParserConfig(cfg.DataParserConfig): + """Instant-NGP dataset parser config""" + + _target: Type = field(default_factory=lambda: InstantNGP) + """target class to instantiate""" + data_directory: 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 + """How much to scale the scene. Defaults to 0.33""" + + @dataclass class InstantNGP(DataParser): """Instant NGP Dataset""" - config: cfg.InstantNGPDataParserConfig + config: InstantNGPDataParserConfig def _generate_dataset_inputs(self, split="train"): diff --git a/nerfactory/datamanagers/dataparsers/mipnerf_parser.py b/nerfactory/datamanagers/dataparsers/mipnerf_parser.py index 58c8cb1ec..2f353a464 100644 --- a/nerfactory/datamanagers/dataparsers/mipnerf_parser.py +++ b/nerfactory/datamanagers/dataparsers/mipnerf_parser.py @@ -15,7 +15,9 @@ """Data parser for mipnerf dataset""" from __future__ import annotations -from dataclasses import dataclass +from dataclasses import dataclass, field +from pathlib import Path +from typing import Type import imageio import numpy as np @@ -28,11 +30,29 @@ from nerfactory.datamanagers.structs import DatasetInputs, SceneBounds from nerfactory.utils.io import get_absolute_path +@dataclass +class MipNerf360DataParserConfig(cfg.DataParserConfig): + """Mipnerf 360 dataset parser config""" + + _target: Type = field(default_factory=lambda: Mipnerf360) + """target class to instantiate""" + data_directory: Path = Path("data/mipnerf_360/garden") + """directory specifying location of data""" + downscale_factor: int = 1 + """How much to downscale images. Defaults to 1.""" + val_skip: int = 8 + """1/val_skip images to use for validation. Defaults to 8.""" + auto_scale: bool = True + """Scale based on pose bounds. Defaults to True.""" + aabb_scale: float = 4 + """Scene scale, Defaults to 1.0.""" + + @dataclass class Mipnerf360(DataParser): """MipNeRF 360 Dataset""" - config: cfg.MipNerf360DataParserConfig + config: MipNerf360DataParserConfig @classmethod def normalize_orientation(cls, poses: np.ndarray): diff --git a/nerfactory/datamanagers/dataparsers/nerfactory_parser.py b/nerfactory/datamanagers/dataparsers/nerfactory_parser.py index 98659d2ae..8a6c564d2 100644 --- a/nerfactory/datamanagers/dataparsers/nerfactory_parser.py +++ b/nerfactory/datamanagers/dataparsers/nerfactory_parser.py @@ -11,13 +11,15 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. +""" Data parser for nerfactory datasets. """ from __future__ import annotations import logging import math -from dataclasses import dataclass +from dataclasses import dataclass, field from pathlib import Path +from typing import Literal, Type import numpy as np import torch @@ -30,11 +32,34 @@ from nerfactory.datamanagers.structs import DatasetInputs, SceneBounds from nerfactory.utils.io import get_absolute_path, load_from_json +@dataclass +class NerfactoryDataParserConfig(cfg.DataParserConfig): + """Nerfactory dataset config""" + + _target: Type = field(default_factory=lambda: Nerfactory) + """target class to instantiate""" + data_directory: Path = Path("data/ours/posters_v3") + """directory specifying location of data""" + scale_factor: float = 1.0 + """How much to scale the camera origins by.""" + downscale_factor: int = 1 + """How much to downscale images. Defaults to 1.""" + scene_scale: float = 4.0 + """How much to scale the scene. Defaults to 4.0""" + orientation_method: Literal["pca", "up"] = "up" + """The method to use for orientation. Either "pca" or "up".""" + train_split_percentage: float = 0.9 + """ + The percent of images to use for training. + The remaining images are for eval. Defaults to 0.9. + """ + + @dataclass class Nerfactory(DataParser): """Nerfactory Dataset""" - config: cfg.NerfactoryDataParserConfig + config: NerfactoryDataParserConfig def _generate_dataset_inputs(self, split="train"): @@ -76,7 +101,7 @@ class Nerfactory(DataParser): assert len(i_eval) == num_eval_images if split == "train": indices = i_train - elif split == "val" or split == "test": + elif split in ["val", "test"]: indices = i_eval else: raise ValueError(f"Unknown dataparser split {split}") diff --git a/nerfactory/datamanagers/dataparsers/record3d_parser.py b/nerfactory/datamanagers/dataparsers/record3d_parser.py index 55598fa98..d70983a6e 100644 --- a/nerfactory/datamanagers/dataparsers/record3d_parser.py +++ b/nerfactory/datamanagers/dataparsers/record3d_parser.py @@ -15,7 +15,9 @@ """Data parser for record3d dataset""" from __future__ import annotations -from dataclasses import dataclass +from dataclasses import dataclass, field +from pathlib import Path +from typing import Literal, Type import numpy as np import torch @@ -30,11 +32,30 @@ from nerfactory.utils import poses as pose_utils from nerfactory.utils.io import get_absolute_path, load_from_json +@dataclass +class Record3DDataParserConfig(cfg.DataParserConfig): + """Record3D dataset config""" + + _target: Type = field(default_factory=lambda: Record3D) + """target class to instantiate""" + data_directory: Path = Path("data/record3d/garden") + """Location of data""" + val_skip: int = 8 + """1/val_skip images to use for validation. Defaults to 8.""" + aabb_scale: float = 4.0 + """Scene scale, Defaults to 4.0.""" + orientation_method: Literal["pca", "up"] = "up" + """The method to use for orientation. Either "pca" or "up".""" + max_dataset_size: int = 150 + """Max number of images to train on. If the dataset has + more, images will be sampled approximately evenly. Defaults to 150.""" + + @dataclass class Record3D(DataParser): """Record3D Dataset""" - config: cfg.Record3DDataParserConfig + config: Record3DDataParserConfig def _generate_dataset_inputs(self, split: str = "train") -> DatasetInputs: abs_dir = get_absolute_path(self.config.data_directory)