diff --git a/docs/developer_guides/new_methods.md b/docs/developer_guides/new_methods.md index 001134b07..6deced451 100644 --- a/docs/developer_guides/new_methods.md +++ b/docs/developer_guides/new_methods.md @@ -78,6 +78,27 @@ Instead, you may use the `NERFSTUDIO_METHOD_CONFIGS` environment variable to tem export NERFSTUDIO_METHOD_CONFIGS="my-method=my_method.my_config:MyMethod" ``` +The `NERFSTUDIO_METHOD_CONFIGS` environment variable additionally accepts a function or derived class to temporarily register your custom method. + +```python +"""my_method/my_config.py""" + +from dataclasses import dataclass +from nerfstudio.engine.trainer import TrainerConfig +from nerfstudio.plugins.types import MethodSpecification + +def MyMethodFunc(): + return = MethodSpecification( + config=TrainerConfig(...) + description="Custom description" + ) + +@dataclass +class MyMethodClass(MethodSpecification): + config: TrainerConfig = TrainerConfig(...) + description: str = "Custom description" +``` + ## 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`. diff --git a/nerfstudio/plugins/registry.py b/nerfstudio/plugins/registry.py index 808772678..f50550f74 100644 --- a/nerfstudio/plugins/registry.py +++ b/nerfstudio/plugins/registry.py @@ -60,7 +60,16 @@ def discover_methods() -> t.Tuple[t.Dict[str, TrainerConfig], t.Dict[str, str]]: CONSOLE.print(f"[bold green]Info: Loading method {name} from environment variable") module, config_name = path.split(":") method_config = getattr(importlib.import_module(module), config_name) - assert isinstance(method_config, MethodSpecification) + + # method_config specified as function or class -> instance + if callable(method_config): + method_config = method_config() + + # check for valid instance type + if not isinstance(method_config, MethodSpecification): + raise TypeError("Method is not an instance of MethodSpecification") + + # save to methods methods[name] = method_config.config descriptions[name] = method_config.description except Exception: diff --git a/tests/plugins/test_registry.py b/tests/plugins/test_registry.py index 1fb993633..08805b2d9 100644 --- a/tests/plugins/test_registry.py +++ b/tests/plugins/test_registry.py @@ -3,6 +3,7 @@ Tests for the nerfstudio.plugins.registry module. """ import os import sys +from dataclasses import dataclass from nerfstudio.engine.trainer import TrainerConfig from nerfstudio.pipelines.base_pipeline import VanillaPipelineConfig @@ -15,14 +16,21 @@ else: from importlib import metadata as importlib_metadata -TestConfig = MethodSpecification( - config=TrainerConfig( +@dataclass +class TestConfigClass(MethodSpecification): + config: TrainerConfig = TrainerConfig( method_name="test-method", pipeline=VanillaPipelineConfig(), optimizers={}, - ), - description="Test description", -) + ) + description: str = "Test description" + + +def TestConfigFunc(): + return TestConfigClass() + + +TestConfig = TestConfigClass() def test_discover_methods(): @@ -52,17 +60,18 @@ def test_discover_methods(): registry.entry_points = entry_points_backup -def test_discover_methods_from_environment_variable(): +def _test_discover_methods_from_environment_variable(method): """Tests if a custom method from env variable gets properly registered using the discover_methods method""" old_env = None try: old_env = os.environ.get("NERFSTUDIO_METHOD_CONFIGS", None) - os.environ["NERFSTUDIO_METHOD_CONFIGS"] = "test-method-env=test_registry:TestConfig" + method_key = f"{method}-env" + os.environ["NERFSTUDIO_METHOD_CONFIGS"] = f"{method_key}=test_registry:{method}" # Discover plugins methods, _ = registry.discover_methods() - assert "test-method-env" in methods - config = methods["test-method-env"] + assert method_key in methods + config = methods[method_key] assert isinstance(config, TrainerConfig) finally: # Revert mock @@ -70,3 +79,18 @@ def test_discover_methods_from_environment_variable(): os.environ["NERFSTUDIO_METHOD_CONFIGS"] = old_env else: del os.environ["NERFSTUDIO_METHOD_CONFIGS"] + + +def test_discover_methods_from_environment_variable_class(): + """Tests if a custom method class from env variable gets properly registered using the discover_methods method""" + _test_discover_methods_from_environment_variable("TestConfigClass") + + +def test_discover_methods_from_environment_variable_function(): + """Tests if a custom method function from env variable gets properly registered using the discover_methods method""" + _test_discover_methods_from_environment_variable("TestConfigFunc") + + +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")