From 0a94a3ecf3156b5fddb25b7a6e92a4b3fdbb30fb Mon Sep 17 00:00:00 2001 From: Jonas Kulhanek Date: Thu, 29 Jun 2023 11:40:11 +0200 Subject: [PATCH] Unify method_config and dataparser_config registration --- nerfstudio/plugins/registry.py | 44 ++++++++++++++++++- nerfstudio/plugins/registry_dataparser.py | 51 ++++++----------------- nerfstudio/plugins/types.py | 14 +++++++ 3 files changed, 69 insertions(+), 40 deletions(-) diff --git a/nerfstudio/plugins/registry.py b/nerfstudio/plugins/registry.py index 808772678..357f0ccee 100644 --- a/nerfstudio/plugins/registry.py +++ b/nerfstudio/plugins/registry.py @@ -21,8 +21,9 @@ import os import sys import typing as t +from nerfstudio.data.dataparsers.base_dataparser import DataParserConfig from nerfstudio.engine.trainer import TrainerConfig -from nerfstudio.plugins.types import MethodSpecification +from nerfstudio.plugins.types import DataParserSpecification, MethodSpecification from nerfstudio.utils.rich_utils import CONSOLE if sys.version_info < (3, 10): @@ -68,3 +69,44 @@ def discover_methods() -> t.Tuple[t.Dict[str, TrainerConfig], t.Dict[str, str]]: CONSOLE.print("[bold red]Error: Could not load methods from environment variable NERFSTUDIO_METHOD_CONFIGS") return methods, descriptions + + +def discover_dataparsers() -> t.Tuple[t.Dict[str, DataParserConfig], t.Dict[str, str]]: + """ + Discovers all DataParsers registered using the `nerfstudio.dataparser_configs` entrypoint. + And also DataParsers in the NERFSTUDIO_DATAPARSER_CONFIGS environment variable. + """ + configs = {} + descriptions = {} + 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 not an instance of DataParserSpecification" + ) + continue + spec = t.cast(DataParserSpecification, spec) + configs[name] = spec.config + descriptions[name] = spec.description + + if "NERFSTUDIO_DATAPARSER_CONFIGS" in os.environ: + try: + strings = os.environ["NERFSTUDIO_DATAPARSER_CONFIGS"].split(",") + for definition in strings: + if not definition: + continue + name, path = definition.split("=") + CONSOLE.print(f"[bold green]Info: Loading method {name} from environment variable") + module, config_name = path.split(":") + dataparser_config = getattr(importlib.import_module(module), config_name) + assert isinstance(dataparser_config, DataParserSpecification) + configs[name] = dataparser_config.config + descriptions[name] = dataparser_config.description + except Exception: + CONSOLE.print_exception() + CONSOLE.print( + "[bold red]Error: Could not load methods from environment variable NERFSTUDIO_DATAPARSER_CONFIGS" + ) + + return configs, descriptions diff --git a/nerfstudio/plugins/registry_dataparser.py b/nerfstudio/plugins/registry_dataparser.py index f15fdc78c..f0c72b7a7 100644 --- a/nerfstudio/plugins/registry_dataparser.py +++ b/nerfstudio/plugins/registry_dataparser.py @@ -13,49 +13,22 @@ # limitations under the License. """ -Module that keeps all registered plugins and allows for plugin discovery. +[LEGACY] 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 +import warnings +from typing import Dict -from nerfstudio.data.dataparsers.base_dataparser import DataParserConfig +from .registry import discover_dataparsers as new_discover_dataparsers +from .types import DataParserConfig # pylint: disable=unused-import -if sys.version_info < (3, 10): - from importlib_metadata import entry_points -else: - from importlib.metadata import entry_points -CONSOLE = Console(width=120) +warnings.warn( + "This module is deprecated and will be removed in future releases. Use nerfstudio.plugins.registry and nerfstudio.plugins.types instead!", + DeprecationWarning, +) -@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 +def discover_dataparsers() -> Dict[str, DataParserConfig]: + configs, _docs = new_discover_dataparsers() + return configs diff --git a/nerfstudio/plugins/types.py b/nerfstudio/plugins/types.py index cfdb54b55..ce8461f05 100644 --- a/nerfstudio/plugins/types.py +++ b/nerfstudio/plugins/types.py @@ -17,6 +17,7 @@ This package contains specifications used to register plugins. """ from dataclasses import dataclass +from nerfstudio.data.dataparsers.base_dataparser import DataParserConfig from nerfstudio.engine.trainer import TrainerConfig @@ -31,3 +32,16 @@ class MethodSpecification: """Trainer configuration""" description: str """Method description shown in `ns-train` help""" + + +@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""" + description: str + """DataParser description shown in `ns-train` help"""