Unify method_config and dataparser_config registration

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

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

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