diff --git a/docs/developer_guides/new_methods.md b/docs/developer_guides/new_methods.md index baafad440..22f7e5efb 100644 --- a/docs/developer_guides/new_methods.md +++ b/docs/developer_guides/new_methods.md @@ -16,6 +16,7 @@ We recommend the following file structure: │ ├── custom_model.py [optional] │ ├── custom_field.py [optional] │ ├── custom_datamanger.py [optional] +│ ├── custom_dataparser.py [optional] │ ├── ... ├── pyproject.toml ``` @@ -76,6 +77,38 @@ Instead, you may use the `NERFSTUDIO_METHOD_CONFIGS` environment variable to tem export NERFSTUDIO_METHOD_CONFIGS="my-method=my_method.my_config:MyMethod" ``` +## Registering custom dataparser with nerfstudio +We also support adding new dataparsers in a similar way. In order to extend the NeRFstudio and register a customized dataparser, you can register it with Nerfstudio as a `nerfstudio.dataparser_configs` entrypoint in the `pyproject.toml` file. Nerfstudio will automatically look for all registered dataparsers and will register them to be used by methods such as `ns-train`. + +You can declare the dataparser in the same config file: +```python +"""my_method/my_config.py""" + +from nerfstudio.plugins.registry_dataparser import DataParserSpecification +from my_method.custom_dataparser import CustomDataparserConfig + +MyDataparser = DataParserSpecification(config=CustomDataparserConfig) +``` + +Then add the following lines in the `pyproject.toml` file, where the entrypoint to the new dataparser is set. + +"""pyproject.toml""" + +[project] +name = "my_method" + +# ... + +[project.entry-points.'nerfstudio.dataparser_configs'] +custom-dataparser = 'my_method.my_config:MyDataparser' +``` + +finally run the following to register the dataparser. + +``` +pip install -e . +``` + ## Running custom method After registering your method you should be able to run the method with, diff --git a/nerfstudio/configs/dataparser_configs.py b/nerfstudio/configs/dataparser_configs.py new file mode 100644 index 000000000..5a3716ff4 --- /dev/null +++ b/nerfstudio/configs/dataparser_configs.py @@ -0,0 +1,67 @@ +# 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. + +""" +Aggregate all the dataparser configs in one location. +""" + +import tyro + +from nerfstudio.data.dataparsers.arkitscenes_dataparser import ( + ARKitScenesDataParserConfig, +) +from nerfstudio.data.dataparsers.blender_dataparser import BlenderDataParserConfig +from nerfstudio.data.dataparsers.dnerf_dataparser import DNeRFDataParserConfig +from nerfstudio.data.dataparsers.dycheck_dataparser import DycheckDataParserConfig +from nerfstudio.data.dataparsers.instant_ngp_dataparser import ( + InstantNGPDataParserConfig, +) +from nerfstudio.data.dataparsers.minimal_dataparser import MinimalDataParserConfig +from nerfstudio.data.dataparsers.nerfosr_dataparser import NeRFOSRDataParserConfig +from nerfstudio.data.dataparsers.nerfstudio_dataparser import NerfstudioDataParserConfig +from nerfstudio.data.dataparsers.nuscenes_dataparser import NuScenesDataParserConfig +from nerfstudio.data.dataparsers.phototourism_dataparser import ( + PhototourismDataParserConfig, +) +from nerfstudio.data.dataparsers.scannet_dataparser import ScanNetDataParserConfig +from nerfstudio.data.dataparsers.sdfstudio_dataparser import SDFStudioDataParserConfig +from nerfstudio.data.dataparsers.sitcoms3d_dataparser import Sitcoms3DDataParserConfig +from nerfstudio.plugins.registry_dataparser import discover_dataparsers + +dataparsers = { + "nerfstudio-data": NerfstudioDataParserConfig(), + "minimal-parser": MinimalDataParserConfig(), + "arkit-data": ARKitScenesDataParserConfig(), + "blender-data": BlenderDataParserConfig(), + "instant-ngp-data": InstantNGPDataParserConfig(), + "nuscenes-data": NuScenesDataParserConfig(), + "dnerf-data": DNeRFDataParserConfig(), + "phototourism-data": PhototourismDataParserConfig(), + "dycheck-data": DycheckDataParserConfig(), + "scannet-data": ScanNetDataParserConfig(), + "sdfstudio-data": SDFStudioDataParserConfig(), + "nerfosr-data": NeRFOSRDataParserConfig(), + "sitcoms3d-data": Sitcoms3DDataParserConfig(), +} + +external_dataparsers = discover_dataparsers() +all_dataparsers = {**dataparsers, **external_dataparsers} +AnnotatedDataParserUnion = tyro.conf.OmitSubcommandPrefixes[ # Omit prefixes of flags in subcommands. + tyro.extras.subcommand_type_from_defaults( + all_dataparsers, + prefix_names=False, # Omit prefixes in subcommands themselves. + ) +] +"""Union over possible dataparser types, annotated with metadata for tyro. This is the +same as the vanilla union, but results in shorter subcommand names.""" diff --git a/nerfstudio/data/datamanagers/base_datamanager.py b/nerfstudio/data/datamanagers/base_datamanager.py index 2be82a395..5069fbc95 100644 --- a/nerfstudio/data/datamanagers/base_datamanager.py +++ b/nerfstudio/data/datamanagers/base_datamanager.py @@ -25,7 +25,6 @@ from pathlib import Path from typing import Any, Dict, Generic, List, Literal, Optional, Tuple, Type, Union import torch -import tyro from torch import nn from torch.nn import Parameter from torch.utils.data import Dataset @@ -36,26 +35,9 @@ from nerfstudio.cameras.camera_optimizers import CameraOptimizerConfig from nerfstudio.cameras.cameras import CameraType from nerfstudio.cameras.rays import RayBundle from nerfstudio.configs.base_config import InstantiateConfig -from nerfstudio.data.dataparsers.arkitscenes_dataparser import ( - ARKitScenesDataParserConfig, -) +from nerfstudio.configs.dataparser_configs import AnnotatedDataParserUnion from nerfstudio.data.dataparsers.base_dataparser import DataparserOutputs from nerfstudio.data.dataparsers.blender_dataparser import BlenderDataParserConfig -from nerfstudio.data.dataparsers.dnerf_dataparser import DNeRFDataParserConfig -from nerfstudio.data.dataparsers.dycheck_dataparser import DycheckDataParserConfig -from nerfstudio.data.dataparsers.instant_ngp_dataparser import ( - InstantNGPDataParserConfig, -) -from nerfstudio.data.dataparsers.minimal_dataparser import MinimalDataParserConfig -from nerfstudio.data.dataparsers.nerfosr_dataparser import NeRFOSRDataParserConfig -from nerfstudio.data.dataparsers.nerfstudio_dataparser import NerfstudioDataParserConfig -from nerfstudio.data.dataparsers.nuscenes_dataparser import NuScenesDataParserConfig -from nerfstudio.data.dataparsers.phototourism_dataparser import ( - PhototourismDataParserConfig, -) -from nerfstudio.data.dataparsers.scannet_dataparser import ScanNetDataParserConfig -from nerfstudio.data.dataparsers.sdfstudio_dataparser import SDFStudioDataParserConfig -from nerfstudio.data.dataparsers.sitcoms3d_dataparser import Sitcoms3DDataParserConfig from nerfstudio.data.datasets.base_dataset import InputDataset from nerfstudio.data.pixel_samplers import ( EquirectangularPixelSampler, @@ -98,30 +80,6 @@ def variable_res_collate(batch: List[Dict]) -> Dict: return new_batch -AnnotatedDataParserUnion = tyro.conf.OmitSubcommandPrefixes[ # Omit prefixes of flags in subcommands. - tyro.extras.subcommand_type_from_defaults( - { - "nerfstudio-data": NerfstudioDataParserConfig(), - "minimal-parser": MinimalDataParserConfig(), - "arkit-data": ARKitScenesDataParserConfig(), - "blender-data": BlenderDataParserConfig(), - "instant-ngp-data": InstantNGPDataParserConfig(), - "nuscenes-data": NuScenesDataParserConfig(), - "dnerf-data": DNeRFDataParserConfig(), - "phototourism-data": PhototourismDataParserConfig(), - "dycheck-data": DycheckDataParserConfig(), - "scannet-data": ScanNetDataParserConfig(), - "sdfstudio-data": SDFStudioDataParserConfig(), - "nerfosr-data": NeRFOSRDataParserConfig(), - "sitcoms3d-data": Sitcoms3DDataParserConfig(), - }, - prefix_names=False, # Omit prefixes in subcommands themselves. - ) -] -"""Union over possible dataparser types, annotated with metadata for tyro. This is the -same as the vanilla union, but results in shorter subcommand names.""" - - @dataclass class DataManagerConfig(InstantiateConfig): """Configuration for data manager instantiation; DataManager is in charge of keeping the train/eval dataparsers; diff --git a/nerfstudio/plugins/registry.py b/nerfstudio/plugins/registry.py index 5654c0775..880170ff9 100644 --- a/nerfstudio/plugins/registry.py +++ b/nerfstudio/plugins/registry.py @@ -40,15 +40,15 @@ def discover_methods() -> t.Tuple[t.Dict[str, TrainerConfig], t.Dict[str, str]]: descriptions = {} discovered_entry_points = entry_points(group="nerfstudio.method_configs") for name in discovered_entry_points.names: - specification = discovered_entry_points[name].load() - if not isinstance(specification, MethodSpecification): + spec = discovered_entry_points[name].load() + if not isinstance(spec, MethodSpecification): CONSOLE.print( - "[bold yellow]Warning: Could not entry point {n} as it is not an instance of MethodSpecification" + f"[bold yellow]Warning: Could not entry point {spec} as it is not an instance of MethodSpecification" ) continue - specification = t.cast(MethodSpecification, specification) - methods[specification.config.method_name] = specification.config - descriptions[specification.config.method_name] = specification.description + spec = t.cast(MethodSpecification, spec) + methods[spec.config.method_name] = spec.config + descriptions[spec.config.method_name] = spec.description if "NERFSTUDIO_METHOD_CONFIGS" in os.environ: try: diff --git a/nerfstudio/plugins/registry_dataparser.py b/nerfstudio/plugins/registry_dataparser.py new file mode 100644 index 000000000..eb96e8751 --- /dev/null +++ b/nerfstudio/plugins/registry_dataparser.py @@ -0,0 +1,61 @@ +# 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. + +""" +Module that keeps all registered plugins and allows for plugin discovery. +""" + +import sys +import typing as t +from dataclasses import dataclass + +from rich.progress import Console + +from nerfstudio.data.dataparsers.base_dataparser import DataParserConfig + +if sys.version_info < (3, 10): + from importlib_metadata import entry_points +else: + from importlib.metadata import entry_points +CONSOLE = Console(width=120) + + +@dataclass +class DataParserSpecification: + """ + DataParser specification class used to register custom dataparsers with Nerfstudio. + The registered dataparsers will be available in commands such as `ns-train` + """ + + config: DataParserConfig + """Dataparser configuration""" + + +def discover_dataparsers() -> t.Dict[str, DataParserConfig]: + """ + Discovers all dataparsers registered using the `nerfstudio.dataparser_configs` entrypoint. + """ + dataparsers = {} + discovered_entry_points = entry_points(group="nerfstudio.dataparser_configs") + for name in discovered_entry_points.names: + spec = discovered_entry_points[name].load() + if not isinstance(spec, DataParserSpecification): + CONSOLE.print( + f"[bold yellow]Warning: Could not entry point {spec} as it is an instance of DataParserSpecification" + ) + continue + spec = t.cast(DataParserSpecification, spec) + dataparsers[name] = spec.config + + return dataparsers