[Enhancement] Support registering function

mmcv/utils/registry.py

@@ -43,7 +43,7 @@ def build_from_cfg(cfg, registry, default_args=None):
         if obj_cls is None:
             raise KeyError(
                 f'{obj_type} is not in the {registry.name} registry')
-    elif inspect.isclass(obj_type):
+    elif inspect.isclass(obj_type) or inspect.isfunction(obj_type):
         obj_cls = obj_type
     else:
         raise TypeError(
@@ -56,16 +56,21 @@ def build_from_cfg(cfg, registry, default_args=None):
 
 
 class Registry:
-    """A registry to map strings to classes.
+    """A registry to map strings to classes or functions.
 
-    Registered object could be built from registry.
+    Registered object could be built from registry. Meanwhile, registered
+    functions could be called from registry.
 
     Example:
         >>> MODELS = Registry('models')
         >>> @MODELS.register_module()
         >>> class ResNet:
         >>>     pass
         >>> resnet = MODELS.build(dict(type='ResNet'))
+        >>> @MODELS.register_module()
+        >>> def resnet50():
+        >>>     pass
+        >>> resnet = MODELS.build(dict(type='resnet50'))
 
     Please refer to
     https://mmcv.readthedocs.io/en/latest/understand_mmcv/registry.html for
@@ -233,8 +238,9 @@ def _add_children(self, registry):
         self.children[registry.scope] = registry
 
     def _register_module(self, module_class, module_name=None, force=False):
-        if not inspect.isclass(module_class):
-            raise TypeError('module must be a class, '
+        if not inspect.isclass(module_class) and not inspect.isfunction(
+                module_class):
+            raise TypeError('module must be a class or a function, '
                             f'but got {type(module_class)}')
 
         if module_name is None:
@@ -286,7 +292,7 @@ def register_module(self, name=None, force=False, module=None):
                 specified, the class name will be used.
             force (bool, optional): Whether to override an existing class with
                 the same name. Default: False.
-            module (type): Module class to be registered.
+            module (type): Module class or function to be registered.
         """
         if not isinstance(force, bool):
             raise TypeError(f'force must be a boolean, but got {type(force)}')

tests/test_utils/test_registry.py

@@ -89,12 +89,23 @@ class SphynxCat:
     with pytest.raises(TypeError):
         CATS.register_module(0)
 
-    # can only decorate a class
+    @CATS.register_module()
+    def muchkin():
+        pass
+
+    assert CATS.get('muchkin') is muchkin
+    assert 'muchkin' in CATS
+
+    # can only decorate a class or a function
     with pytest.raises(TypeError):
 
-        @CATS.register_module()
-        def some_method():
-            pass
+        class Demo:
+
+            def some_method(self):
+                pass
+
+        method = Demo().some_method
+        CATS.register_module(name='some_method', module=method)
 
     # begin: test old APIs
     with pytest.warns(DeprecationWarning):