sczhou/CodeFormer

1# Modified from: https://github.com/facebookresearch/fvcore/blob/master/fvcore/common/registry.py  # noqa: E5012 3 4class Registry():5    """6    The registry that provides name -> object mapping, to support third-party7    users' custom modules.8 9    To create a registry (e.g. a backbone registry):10 11    .. code-block:: python12 13        BACKBONE_REGISTRY = Registry('BACKBONE')14 15    To register an object:16 17    .. code-block:: python18 19        @BACKBONE_REGISTRY.register()20        class MyBackbone():21            ...22 23    Or:24 25    .. code-block:: python26 27        BACKBONE_REGISTRY.register(MyBackbone)28    """29 30    def __init__(self, name):31        """32        Args:33            name (str): the name of this registry34        """35        self._name = name36        self._obj_map = {}37 38    def _do_register(self, name, obj):39        assert (name not in self._obj_map), (f"An object named '{name}' was already registered "40                                             f"in '{self._name}' registry!")41        self._obj_map[name] = obj42 43    def register(self, obj=None):44        """45        Register the given object under the the name `obj.__name__`.46        Can be used as either a decorator or not.47        See docstring of this class for usage.48        """49        if obj is None:50            # used as a decorator51            def deco(func_or_class):52                name = func_or_class.__name__53                self._do_register(name, func_or_class)54                return func_or_class55 56            return deco57 58        # used as a function call59        name = obj.__name__60        self._do_register(name, obj)61 62    def get(self, name):63        ret = self._obj_map.get(name)64        if ret is None:65            raise KeyError(f"No object named '{name}' found in '{self._name}' registry!")66        return ret67 68    def __contains__(self, name):69        return name in self._obj_map70 71    def __iter__(self):72        return iter(self._obj_map.items())73 74    def keys(self):75        return self._obj_map.keys()76 77 78DATASET_REGISTRY = Registry('dataset')79ARCH_REGISTRY = Registry('arch')80MODEL_REGISTRY = Registry('model')81LOSS_REGISTRY = Registry('loss')82METRIC_REGISTRY = Registry('metric')83