Unify dataparser and method discovery (#2398)

* Unify external dataparser registration

* Unify dataparser and method discovery

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,

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:

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

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")