From 1147fe4d5c360995774a77bd282f5de7dc167779 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jon=C3=A1=C5=A1=20Kulh=C3=A1nek?= Date: Sat, 6 May 2023 18:54:52 +0200 Subject: [PATCH] Make VanillaDataManager generic (#1868) * Make VanillaDataManager generic * Remove removed classes from docs * Simplify generic datamanager * Simplify code * Add serialization tests --- docs/reference/api/data/datamanagers.rst | 23 +--- nerfstudio/configs/method_configs.py | 36 ++++-- .../data/datamanagers/base_datamanager.py | 32 ++++-- .../data/datamanagers/depth_datamanager.py | 48 -------- .../data/datamanagers/sdf_datamanager.py | 48 -------- .../data/datamanagers/semantic_datamanager.py | 54 --------- tests/data/test_datamanager.py | 107 ++++++++++++++++++ 7 files changed, 155 insertions(+), 193 deletions(-) delete mode 100644 nerfstudio/data/datamanagers/depth_datamanager.py delete mode 100644 nerfstudio/data/datamanagers/sdf_datamanager.py delete mode 100644 nerfstudio/data/datamanagers/semantic_datamanager.py create mode 100644 tests/data/test_datamanager.py diff --git a/docs/reference/api/data/datamanagers.rst b/docs/reference/api/data/datamanagers.rst index 7779c9e10..43e5d6cce 100644 --- a/docs/reference/api/data/datamanagers.rst +++ b/docs/reference/api/data/datamanagers.rst @@ -7,25 +7,4 @@ Base ---------------- .. automodule:: nerfstudio.data.datamanagers.base_datamanager - :members: - -Depth Datamanager ------------------ - -.. automodule:: nerfstudio.data.datamanagers.depth_datamanager - :members: - :show-inheritance: - -SDF Datamanager ----------------- - -.. automodule:: nerfstudio.data.datamanagers.sdf_datamanager - :members: - :show-inheritance: - -Semantic Datamanager ---------------------- - -.. automodule:: nerfstudio.data.datamanagers.semantic_datamanager - :members: - :show-inheritance: \ No newline at end of file + :members: \ No newline at end of file diff --git a/nerfstudio/configs/method_configs.py b/nerfstudio/configs/method_configs.py index 909c5c388..3390a1bfa 100644 --- a/nerfstudio/configs/method_configs.py +++ b/nerfstudio/configs/method_configs.py @@ -24,10 +24,10 @@ import tyro from nerfstudio.cameras.camera_optimizers import CameraOptimizerConfig from nerfstudio.configs.base_config import ViewerConfig -from nerfstudio.data.datamanagers.base_datamanager import VanillaDataManagerConfig -from nerfstudio.data.datamanagers.depth_datamanager import DepthDataManagerConfig -from nerfstudio.data.datamanagers.sdf_datamanager import SDFDataManagerConfig -from nerfstudio.data.datamanagers.semantic_datamanager import SemanticDataManagerConfig +from nerfstudio.data.datamanagers.base_datamanager import ( + VanillaDataManager, + VanillaDataManagerConfig, +) from nerfstudio.data.dataparsers.blender_dataparser import BlenderDataParserConfig from nerfstudio.data.dataparsers.dnerf_dataparser import DNeRFDataParserConfig from nerfstudio.data.dataparsers.dycheck_dataparser import DycheckDataParserConfig @@ -40,6 +40,9 @@ from nerfstudio.data.dataparsers.phototourism_dataparser import ( ) from nerfstudio.data.dataparsers.sdfstudio_dataparser import SDFStudioDataParserConfig from nerfstudio.data.dataparsers.sitcoms3d_dataparser import Sitcoms3DDataParserConfig +from nerfstudio.data.datasets.depth_dataset import DepthDataset +from nerfstudio.data.datasets.sdf_dataset import SDFDataset +from nerfstudio.data.datasets.semantic_dataset import SemanticDataset from nerfstudio.engine.optimizers import AdamOptimizerConfig, RAdamOptimizerConfig from nerfstudio.engine.schedulers import ( CosineDecaySchedulerConfig, @@ -162,7 +165,8 @@ method_configs["depth-nerfacto"] = TrainerConfig( max_num_iterations=30000, mixed_precision=True, pipeline=VanillaPipelineConfig( - datamanager=DepthDataManagerConfig( + datamanager=VanillaDataManagerConfig( + _target=VanillaDataManager[DepthDataset], dataparser=NerfstudioDataParserConfig(), train_num_rays_per_batch=4096, eval_num_rays_per_batch=4096, @@ -308,8 +312,11 @@ method_configs["semantic-nerfw"] = TrainerConfig( max_num_iterations=30000, mixed_precision=True, pipeline=VanillaPipelineConfig( - datamanager=SemanticDataManagerConfig( - dataparser=Sitcoms3DDataParserConfig(), train_num_rays_per_batch=4096, eval_num_rays_per_batch=8192 + datamanager=VanillaDataManagerConfig( + _target=VanillaDataManager[SemanticDataset], + dataparser=Sitcoms3DDataParserConfig(), + train_num_rays_per_batch=4096, + eval_num_rays_per_batch=8192, ), model=SemanticNerfWModelConfig(eval_num_rays_per_chunk=1 << 16), ), @@ -440,7 +447,8 @@ method_configs["nerfplayer-nerfacto"] = TrainerConfig( max_num_iterations=30000, mixed_precision=True, pipeline=VanillaPipelineConfig( - datamanager=DepthDataManagerConfig( + datamanager=VanillaDataManagerConfig( + _target=VanillaDataManager[DepthDataset], dataparser=DycheckDataParserConfig(), train_num_rays_per_batch=4096, eval_num_rays_per_batch=4096, @@ -471,7 +479,11 @@ method_configs["nerfplayer-ngp"] = TrainerConfig( max_num_iterations=30000, mixed_precision=True, pipeline=DynamicBatchPipelineConfig( - datamanager=DepthDataManagerConfig(dataparser=DycheckDataParserConfig(), train_num_rays_per_batch=8192), + datamanager=VanillaDataManagerConfig( + _target=VanillaDataManager[DepthDataset], + dataparser=DycheckDataParserConfig(), + train_num_rays_per_batch=8192, + ), model=NerfplayerNGPModelConfig( eval_num_rays_per_chunk=8192, grid_levels=1, @@ -500,7 +512,8 @@ method_configs["neus"] = TrainerConfig( max_num_iterations=100000, mixed_precision=False, pipeline=VanillaPipelineConfig( - datamanager=SDFDataManagerConfig( + datamanager=VanillaDataManagerConfig( + _target=VanillaDataManager[SDFDataset], dataparser=SDFStudioDataParserConfig(), train_num_rays_per_batch=1024, eval_num_rays_per_batch=1024, @@ -533,7 +546,8 @@ method_configs["neus-facto"] = TrainerConfig( max_num_iterations=20001, mixed_precision=False, pipeline=VanillaPipelineConfig( - datamanager=SDFDataManagerConfig( + datamanager=VanillaDataManagerConfig( + _target=VanillaDataManager[SDFDataset], dataparser=SDFStudioDataParserConfig(), train_num_rays_per_batch=2048, eval_num_rays_per_batch=2048, diff --git a/nerfstudio/data/datamanagers/base_datamanager.py b/nerfstudio/data/datamanagers/base_datamanager.py index d696519af..0c495139e 100644 --- a/nerfstudio/data/datamanagers/base_datamanager.py +++ b/nerfstudio/data/datamanagers/base_datamanager.py @@ -18,10 +18,11 @@ Datamanager. from __future__ import annotations +import functools from abc import abstractmethod from dataclasses import dataclass, field from pathlib import Path -from typing import Any, Dict, List, Optional, Tuple, Type, Union +from typing import Any, Dict, Generic, List, Optional, Tuple, Type, Union import torch import tyro @@ -30,7 +31,7 @@ from torch import nn from torch.nn import Parameter from torch.utils.data import Dataset from torch.utils.data.distributed import DistributedSampler -from typing_extensions import Literal +from typing_extensions import Literal, TypeVar from nerfstudio.cameras.camera_optimizers import CameraOptimizerConfig from nerfstudio.cameras.cameras import CameraType @@ -367,7 +368,10 @@ class VanillaDataManagerConfig(DataManagerConfig): """Size of patch to sample from. If >1, patch-based sampling will be used.""" -class VanillaDataManager(DataManager): # pylint: disable=abstract-method +TDataset = TypeVar("TDataset", bound=InputDataset, default=InputDataset) + + +class VanillaDataManager(DataManager, Generic[TDataset]): # pylint: disable=abstract-method """Basic stored data manager implementation. This is pretty much a port over from our old dataloading utilities, and is a little jank @@ -381,8 +385,8 @@ class VanillaDataManager(DataManager): # pylint: disable=abstract-method """ config: VanillaDataManagerConfig - train_dataset: InputDataset - eval_dataset: InputDataset + train_dataset: TDataset + eval_dataset: TDataset train_dataparser_outputs: DataparserOutputs train_pixel_sampler: Optional[PixelSampler] = None eval_pixel_sampler: Optional[PixelSampler] = None @@ -396,6 +400,7 @@ class VanillaDataManager(DataManager): # pylint: disable=abstract-method local_rank: int = 0, **kwargs, # pylint: disable=unused-argument ): + self.dataset_type: Type[TDataset] = kwargs.get("_dataset_type", TDataset.__default__) self.config = config self.device = device self.world_size = world_size @@ -426,22 +431,29 @@ class VanillaDataManager(DataManager): # pylint: disable=abstract-method super().__init__() - def create_train_dataset(self) -> InputDataset: + def __class_getitem__(cls, item): + return type( + cls.__name__, + (cls,), + {"__module__": cls.__module__, "__init__": functools.partialmethod(cls.__init__, _dataset_type=item)}, + ) + + def create_train_dataset(self) -> TDataset: """Sets up the data loaders for training""" - return InputDataset( + return self.dataset_type( dataparser_outputs=self.train_dataparser_outputs, scale_factor=self.config.camera_res_scale_factor, ) - def create_eval_dataset(self) -> InputDataset: + def create_eval_dataset(self) -> TDataset: """Sets up the data loaders for evaluation""" - return InputDataset( + return self.dataset_type( dataparser_outputs=self.dataparser.get_dataparser_outputs(split=self.test_split), scale_factor=self.config.camera_res_scale_factor, ) def _get_pixel_sampler( # pylint: disable=no-self-use - self, dataset: InputDataset, *args: Any, **kwargs: Any + self, dataset: TDataset, *args: Any, **kwargs: Any ) -> PixelSampler: """Infer pixel sampler to use.""" if self.config.patch_size > 1: diff --git a/nerfstudio/data/datamanagers/depth_datamanager.py b/nerfstudio/data/datamanagers/depth_datamanager.py deleted file mode 100644 index 465e79e2d..000000000 --- a/nerfstudio/data/datamanagers/depth_datamanager.py +++ /dev/null @@ -1,48 +0,0 @@ -# Copyright 2022 The Nerfstudio 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. - -""" -Depth datamanager. -""" - -from dataclasses import dataclass, field -from typing import Type - -from nerfstudio.data.datamanagers import base_datamanager -from nerfstudio.data.datasets.depth_dataset import DepthDataset - - -@dataclass -class DepthDataManagerConfig(base_datamanager.VanillaDataManagerConfig): - """A depth datamanager - required to use with .setup()""" - - _target: Type = field(default_factory=lambda: DepthDataManager) - - -class DepthDataManager(base_datamanager.VanillaDataManager): # pylint: disable=abstract-method - """Data manager implementation for data that also requires processing depth data. - Args: - config: the DataManagerConfig used to instantiate class - """ - - def create_train_dataset(self) -> DepthDataset: - self.train_dataparser_outputs = self.dataparser.get_dataparser_outputs(split="train") - return DepthDataset( - dataparser_outputs=self.train_dataparser_outputs, - ) - - def create_eval_dataset(self) -> DepthDataset: - return DepthDataset( - dataparser_outputs=self.dataparser.get_dataparser_outputs(split=self.test_split), - ) diff --git a/nerfstudio/data/datamanagers/sdf_datamanager.py b/nerfstudio/data/datamanagers/sdf_datamanager.py deleted file mode 100644 index e1aca46df..000000000 --- a/nerfstudio/data/datamanagers/sdf_datamanager.py +++ /dev/null @@ -1,48 +0,0 @@ -# Copyright 2022 The Nerfstudio 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. - -""" -SDF datamanager. -""" - -from dataclasses import dataclass, field -from typing import Type - -from nerfstudio.data.datamanagers import base_datamanager -from nerfstudio.data.datasets.sdf_dataset import SDFDataset - - -@dataclass -class SDFDataManagerConfig(base_datamanager.VanillaDataManagerConfig): - """A sdf datamanager - required to use with .setup()""" - - _target: Type = field(default_factory=lambda: SDFDataManager) - - -class SDFDataManager(base_datamanager.VanillaDataManager): # pylint: disable=abstract-method - """Data manager implementation for data that also requires processing sdf data. - Args: - config: the DataManagerConfig used to instantiate class - """ - - def create_train_dataset(self) -> SDFDataset: - self.train_dataparser_outputs = self.dataparser.get_dataparser_outputs(split="train") - return SDFDataset( - dataparser_outputs=self.train_dataparser_outputs, - ) - - def create_eval_dataset(self) -> SDFDataset: - return SDFDataset( - dataparser_outputs=self.dataparser.get_dataparser_outputs(split=self.test_split), - ) diff --git a/nerfstudio/data/datamanagers/semantic_datamanager.py b/nerfstudio/data/datamanagers/semantic_datamanager.py deleted file mode 100644 index 322b1d3bd..000000000 --- a/nerfstudio/data/datamanagers/semantic_datamanager.py +++ /dev/null @@ -1,54 +0,0 @@ -# Copyright 2022 The Nerfstudio 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. - -""" -Semantic datamanager. -""" - -from dataclasses import dataclass, field -from typing import Type - -from nerfstudio.data.datamanagers.base_datamanager import ( - VanillaDataManager, - VanillaDataManagerConfig, -) -from nerfstudio.data.datasets.semantic_dataset import SemanticDataset - - -@dataclass -class SemanticDataManagerConfig(VanillaDataManagerConfig): - """A semantic datamanager - required to use with .setup()""" - - _target: Type = field(default_factory=lambda: SemanticDataManager) - - -class SemanticDataManager(VanillaDataManager): # pylint: disable=abstract-method - """Data manager implementation for data that also requires processing semantic data. - - Args: - config: the DataManagerConfig used to instantiate class - """ - - def create_train_dataset(self) -> SemanticDataset: - self.train_dataparser_outputs = self.dataparser.get_dataparser_outputs(split="train") - return SemanticDataset( - dataparser_outputs=self.train_dataparser_outputs, - scale_factor=self.config.camera_res_scale_factor, - ) - - def create_eval_dataset(self) -> SemanticDataset: - return SemanticDataset( - dataparser_outputs=self.dataparser.get_dataparser_outputs(split=self.test_split), - scale_factor=self.config.camera_res_scale_factor, - ) diff --git a/tests/data/test_datamanager.py b/tests/data/test_datamanager.py new file mode 100644 index 000000000..e80330183 --- /dev/null +++ b/tests/data/test_datamanager.py @@ -0,0 +1,107 @@ +# pylint: disable=all +from copy import copy +from typing import Any + +import pytest +import torch +import yaml + +from nerfstudio.cameras.cameras import Cameras +from nerfstudio.configs.base_config import InstantiateConfig +from nerfstudio.data.datamanagers.base_datamanager import ( + DataManager, + DataparserOutputs, + VanillaDataManager, + VanillaDataManagerConfig, +) +from nerfstudio.data.datasets.base_dataset import InputDataset +from nerfstudio.data.datasets.depth_dataset import DepthDataset + + +class DummyDataParser: + def __init__(self, *args, **kwargs): + pass + + def __getattr__(self, __name: str) -> Any: + if __name.startswith("_"): + return object.__getattribute__(self, __name) + return None + + def get_dataparser_outputs(self, *args, **kwargs): + return DataparserOutputs( + [], + Cameras( + torch.ones((0, 3, 4)), + torch.ones((0, 1)), + torch.ones((0, 1)), + torch.ones((0, 1)), + torch.ones((0, 1)), + 10, + 10, + ), + metadata={"depth_filenames": [], "depth_unit_scale_factor": 1.0}, + ) + + +@pytest.fixture +def config(): + config = VanillaDataManagerConfig() + setattr(config, "dataparser", InstantiateConfig(_target=DummyDataParser)) + setattr(config.dataparser, "data", None) + return config + + +def test_data_manager_type_inference(config): + # Mock for a faster test + + assert VanillaDataManager(config).dataset_type is InputDataset + assert VanillaDataManager[DepthDataset](config).dataset_type is DepthDataset + + class tmp(VanillaDataManager): + pass + + assert tmp(config).dataset_type is InputDataset + + class tmp2(VanillaDataManager[DepthDataset]): + pass + + assert tmp2(config).dataset_type is DepthDataset + + +def test_data_manager_type_can_be_serialized(config): + # Mock for a faster test + + assert VanillaDataManager(config).dataset_type is InputDataset + obj = yaml.load(yaml.dump(VanillaDataManager(config)), Loader=yaml.UnsafeLoader) + assert obj.dataset_type is InputDataset + assert isinstance(obj, VanillaDataManager) + + assert VanillaDataManager[DepthDataset](config).dataset_type is DepthDataset + obj = yaml.load(yaml.dump(VanillaDataManager[DepthDataset](config)), Loader=yaml.UnsafeLoader) + assert obj.dataset_type is DepthDataset + assert isinstance(obj, VanillaDataManager) + + class tmp(VanillaDataManager): + pass + + try: + globals()["tmp"] = tmp + assert tmp(config).dataset_type is InputDataset + obj = yaml.load(yaml.dump(tmp(config)), Loader=yaml.UnsafeLoader) + assert obj.dataset_type is InputDataset + assert isinstance(obj, tmp) + finally: + globals().pop("tmp") + + class tmp2(VanillaDataManager[DepthDataset]): + pass + + try: + globals()["tmp2"] = tmp2 + + assert tmp2(config).dataset_type is DepthDataset + obj = yaml.load(yaml.dump(tmp2(config)), Loader=yaml.UnsafeLoader) + assert obj.dataset_type is DepthDataset + assert isinstance(obj, tmp2) + finally: + globals().pop("tmp2")