diff --git a/docs/developer_guides/new_methods.md b/docs/developer_guides/new_methods.md index a8f8bd423..4b56d8581 100644 --- a/docs/developer_guides/new_methods.md +++ b/docs/developer_guides/new_methods.md @@ -132,6 +132,15 @@ finally run the following to register the dataparser. pip install -e . ``` +Similarly to the method develomement, you can also use environment variables to register dataparsers. +Use the `NERFSTUDIO_DATAPARSER_CONFIGS` environment variable: + +``` +export NERFSTUDIO_DATAPARSER_CONFIGS="my-dataparser=my_package.my_config:MyDataParser" +``` + +Same as with custom methods, `NERFSTUDIO_DATAPARSER_CONFIGS` environment variable additionally accepts a function or derived class to temporarily register your custom method. + ## 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 index 46a728141..f60092a46 100644 --- a/nerfstudio/configs/dataparser_configs.py +++ b/nerfstudio/configs/dataparser_configs.py @@ -54,7 +54,7 @@ dataparsers = { "colmap": ColmapDataParserConfig(), } -external_dataparsers = discover_dataparsers() +external_dataparsers, _ = discover_dataparsers() all_dataparsers = {**dataparsers, **external_dataparsers} if TYPE_CHECKING: diff --git a/nerfstudio/plugins/registry_dataparser.py b/nerfstudio/plugins/registry_dataparser.py index f15fdc78c..e35aeaf48 100644 --- a/nerfstudio/plugins/registry_dataparser.py +++ b/nerfstudio/plugins/registry_dataparser.py @@ -16,6 +16,8 @@ Module that keeps all registered plugins and allows for plugin discovery. """ +import os +import importlib import sys import typing as t from dataclasses import dataclass @@ -41,21 +43,55 @@ class DataParserSpecification: config: DataParserConfig """Dataparser configuration""" + description: t.Optional[str] = None + """Description of the dataparser""" -def discover_dataparsers() -> t.Dict[str, DataParserConfig]: + +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. """ dataparsers = {} + 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 an instance of DataParserSpecification" + f"[bold yellow]Warning: Could not entry point {spec} as it is not an instance of DataParserSpecification" ) continue spec = t.cast(DataParserSpecification, spec) dataparsers[name] = spec.config + descriptions[name] = spec.description - return dataparsers + 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) + + # method_config specified as function or class -> instance + if callable(dataparser_config): + dataparser_config = dataparser_config() + + # check for valid instance type + if not isinstance(dataparser_config, DataParserSpecification): + raise TypeError("Method is not an instance of DataParserSpecification") + + # save to methods + dataparsers[name] = dataparser_config.config + descriptions[name] = dataparser_config.description + except Exception: # pylint: disable=broad-except + CONSOLE.print_exception() + CONSOLE.print( + "[bold red]Error: Could not load methods from environment variable NERFSTUDIO_DATAPARSER_CONFIGS" + ) + + return dataparsers, descriptions diff --git a/tests/plugins/test_registry.py b/tests/plugins/test_registry.py index 08805b2d9..7263a79c0 100644 --- a/tests/plugins/test_registry.py +++ b/tests/plugins/test_registry.py @@ -9,6 +9,8 @@ from nerfstudio.engine.trainer import TrainerConfig from nerfstudio.pipelines.base_pipeline import VanillaPipelineConfig from nerfstudio.plugins import registry from nerfstudio.plugins.types import MethodSpecification +from nerfstudio.plugins import registry_dataparser +from nerfstudio.plugins.registry_dataparser import DataParserSpecification, discover_dataparsers, DataParserConfig if sys.version_info < (3, 10): import importlib_metadata @@ -94,3 +96,79 @@ def test_discover_methods_from_environment_variable_function(): def test_discover_methods_from_environment_variable_instance(): """Tests if a custom method instance from env variable gets properly registered using the discover_methods method""" _test_discover_methods_from_environment_variable("TestConfig") + + +@dataclass +class TestDataparserConfigClass(DataParserSpecification): + config: DataParserConfig = DataParserConfig() + description: str = "Test description" + + +def TestDataparserConfigFunc(): + return TestDataparserConfigClass() + + +TestDataparserConfig = TestDataparserConfigClass() + + +def test_discover_dataparser(): + """Tests if a custom method gets properly registered using the discover_dataparsers method""" + entry_points_backup = registry_dataparser.entry_points + + def entry_points(group=None): + return importlib_metadata.EntryPoints( + [ + importlib_metadata.EntryPoint( + name="test", value="test_registry:TestDataparserConfig", group="nerfstudio.dataparser_configs" + ) + ] + ).select(group=group) + + try: + # Mock importlib entry_points + registry_dataparser.entry_points = entry_points + + # Discover plugins + dataparsers, _ = discover_dataparsers() + assert "test" in dataparsers + config = dataparsers["test"] + assert isinstance(config, DataParserConfig) + finally: + # Revert mock + registry_dataparser.entry_points = entry_points_backup + + +def _test_discover_dataparsers_from_environment_variable(method): + """Tests if a custom method from env variable gets properly registered using the discover_dataparsers method""" + old_env = None + try: + old_env = os.environ.get("NERFSTUDIO_DATAPARSER_CONFIGS", None) + method_key = f"{method}-env" + os.environ["NERFSTUDIO_DATAPARSER_CONFIGS"] = f"{method_key}=test_registry:{method}" + + # Discover plugins + dataparsers, _ = discover_dataparsers() + assert method_key in dataparsers + config = dataparsers[method_key] + assert isinstance(config, DataParserConfig) + finally: + # Revert mock + if old_env is not None: + os.environ["NERFSTUDIO_DATAPARSER_CONFIGS"] = old_env + else: + del os.environ["NERFSTUDIO_DATAPARSER_CONFIGS"] + + +def test_discover_dataparsers_from_environment_variable_class(): + """Tests if a custom dataparser class from env variable gets properly registered using the discover_methods method""" + _test_discover_dataparsers_from_environment_variable("TestDataparserConfigClass") + + +def test_discover_dataparsers_from_environment_variable_function(): + """Tests if a custom dataparser function from env variable gets properly registered using the discover_methods method""" + _test_discover_dataparsers_from_environment_variable("TestDataparserConfigFunc") + + +def test_discover_dataparsers_from_environment_variable_instance(): + """Tests if a custom dataparser instance from env variable gets properly registered using the discover_methods method""" + _test_discover_dataparsers_from_environment_variable("TestDataparserConfig")