Skip to content

Commit

Permalink
additional registry fixes
Browse files Browse the repository at this point in the history
  • Loading branch information
kozlov721 committed Jan 16, 2025
1 parent 466baab commit ecea701
Show file tree
Hide file tree
Showing 3 changed files with 10 additions and 11 deletions.
4 changes: 2 additions & 2 deletions luxonis_ml/data/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -22,14 +22,14 @@ def load_dataset_plugins() -> None: # pragma: no cover
"""Registers any external dataset BaseDataset class plugins."""
for entry_point in pkg_resources.iter_entry_points("dataset_plugins"):
plugin_class = entry_point.load()
DATASETS_REGISTRY.register_module(module=plugin_class)
DATASETS_REGISTRY.register(module=plugin_class)


def load_loader_plugins() -> None: # pragma: no cover
"""Registers any external dataset BaseLoader class plugins."""
for entry_point in pkg_resources.iter_entry_points("loader_plugins"):
plugin_class = entry_point.load()
DATASETS_REGISTRY.register_module(module=plugin_class)
DATASETS_REGISTRY.register(module=plugin_class)


load_dataset_plugins()
Expand Down
6 changes: 3 additions & 3 deletions luxonis_ml/data/augmentations/custom/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,8 +12,8 @@
"albumentations_transformations"
)

TRANSFORMATIONS.register_module(module=LetterboxResize)
TRANSFORMATIONS.register_module(module=MixUp)
TRANSFORMATIONS.register_module(module=Mosaic4)
TRANSFORMATIONS.register(module=LetterboxResize)
TRANSFORMATIONS.register(module=MixUp)
TRANSFORMATIONS.register(module=Mosaic4)

__all__ = ["LetterboxResize", "MixUp", "Mosaic4", "TRANSFORMATIONS"]
11 changes: 5 additions & 6 deletions luxonis_ml/utils/registry.py
Original file line number Diff line number Diff line change
Expand Up @@ -76,7 +76,8 @@ def register_module(
force: bool = False,
) -> Union[T, Callable[[T], T]]:
warnings.warn(
"`register_module` is deprecated, use `register` instead."
"`register_module` is deprecated, use `register` instead.",
stacklevel=2,
)

return self.register(name=name, module=module, force=force)
Expand Down Expand Up @@ -111,14 +112,14 @@ def register(
Can be used as a decorator or as a normal method:
>>> registry = Registry(name="modules")
>>> @registry.register_module()
>>> @registry.register()
... class Foo:
... pass
>>> registry.get("Foo")
<class '__main__.Foo'>
>>> class Bar:
... pass
>>> registry.register_module(module=Bar)
>>> registry.register(module=Bar)
>>> registry.get("Bar")
<class '__main__.Bar'>
Expand Down Expand Up @@ -227,7 +228,5 @@ def __new__(
)
if register:
registry = registry if registry is not None else new_class.REGISTRY
registry.register_module(
name=register_name or name, module=new_class
)
registry.register(name=register_name or name, module=new_class)
return new_class

0 comments on commit ecea701

Please sign in to comment.