extract dataparser configs to a stand-alone file, support custom dataparsers (#1909)

* extract dataparser configs to a stand-alone file, support custom dataparsers

* add license

* fix pylint issues

* fix type error issue

* update naming; add documentations;

* print f string

* fix pylint issue

---------

Co-authored-by: Matthew Tancik <tancik@berkeley.edu>
This commit is contained in:
wuziruiandMatthew Tancik authored and GitHub committed 2023-05-12 15:11:46 -07:00
1 parent fe4586839a
commit bc110d7bd9
5 files changed
+168 -49

No files matched your search

+6 -6
View File
@@ -40,15 +40,15 @@ def discover_methods() -> t.Tuple[t.Dict[str, TrainerConfig], t.Dict[str, str]]:
descriptions = {}
discovered_entry_points = entry_points(group="nerfstudio.method_configs")
for name in discovered_entry_points.names:
specification = discovered_entry_points[name].load()
if not isinstance(specification, MethodSpecification):
spec = discovered_entry_points[name].load()
if not isinstance(spec, MethodSpecification):
CONSOLE.print(
"[bold yellow]Warning: Could not entry point {n} as it is not an instance of MethodSpecification"
f"[bold yellow]Warning: Could not entry point {spec} as it is not an instance of MethodSpecification"
)
continue
specification = t.cast(MethodSpecification, specification)
methods[specification.config.method_name] = specification.config
descriptions[specification.config.method_name] = specification.description
spec = t.cast(MethodSpecification, spec)
methods[spec.config.method_name] = spec.config
descriptions[spec.config.method_name] = spec.description
if "NERFSTUDIO_METHOD_CONFIGS" in os.environ:
try:
+61
View File
@@ -0,0 +1,61 @@
# Copyright 2022 The Nerfstudio Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
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
from nerfstudio.data.dataparsers.base_dataparser import DataParserConfig
if sys.version_info < (3, 10):
from importlib_metadata import entry_points
else:
from importlib.metadata import entry_points
CONSOLE = Console(width=120)
@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