Move dataset parser configs to files (#476)

* Move parser configs to files

* fixed circular import errors

Co-authored-by: Ethan Weber <eweb0124@gmail.com>

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(
     {

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.

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

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

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)
 

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"):
 

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):

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}")

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)