Add sampler and assigner registry (open-mmlab#2419)

* add sampler and assigner registry

* rename with bbox prefix

* restore __init__

* roll back atss_head

* change import level

* import from sampler/assigner

mmdet/core/bbox/__init__.py

@@ -8,7 +8,7 @@
                          bbox_mapping, bbox_mapping_back, delta2bbox,
                          distance2bbox, roi2bbox)
 
-from .assign_sampling import (  # isort:skip, avoid recursive imports
+from .builder import (  # isort:skip, avoid recursive imports
     assign_and_sample, build_assigner, build_sampler)
 
 __all__ = [

mmdet/core/bbox/assigners/approx_max_iou_assigner.py

@@ -1,9 +1,11 @@
 import torch
 
 from ..geometry import bbox_overlaps
+from ..registry import BBOX_ASSIGNERS
 from .max_iou_assigner import MaxIoUAssigner
 
 
+@BBOX_ASSIGNERS.register_module
 class ApproxMaxIoUAssigner(MaxIoUAssigner):
     """Assign a corresponding gt bbox or background to each bbox.
 

mmdet/core/bbox/assigners/atss_assigner.py

@@ -1,10 +1,12 @@
 import torch
 
 from ..geometry import bbox_overlaps
+from ..registry import BBOX_ASSIGNERS
 from .assign_result import AssignResult
 from .base_assigner import BaseAssigner
 
 
+@BBOX_ASSIGNERS.register_module
 class ATSSAssigner(BaseAssigner):
     """Assign a corresponding gt bbox or background to each bbox.
 

mmdet/core/bbox/assigners/max_iou_assigner.py

@@ -1,10 +1,12 @@
 import torch
 
 from ..geometry import bbox_overlaps
+from ..registry import BBOX_ASSIGNERS
 from .assign_result import AssignResult
 from .base_assigner import BaseAssigner
 
 
+@BBOX_ASSIGNERS.register_module
 class MaxIoUAssigner(BaseAssigner):
     """Assign a corresponding gt bbox or background to each bbox.
 

mmdet/core/bbox/assigners/point_assigner.py

@@ -1,9 +1,11 @@
 import torch
 
+from ..registry import BBOX_ASSIGNERS
 from .assign_result import AssignResult
 from .base_assigner import BaseAssigner
 
 
+@BBOX_ASSIGNERS.register_module
 class PointAssigner(BaseAssigner):
     """Assign a corresponding gt bbox or background to each point.
 

mmdet/core/bbox/builder.py

@@ -0,0 +1,27 @@
+from mmdet.utils import build_from_cfg
+from .assigners import BaseAssigner
+from .registry import BBOX_ASSIGNERS, BBOX_SAMPLERS
+from .samplers import BaseSampler
+
+
+def build_assigner(cfg, **default_args):
+    if isinstance(cfg, BaseAssigner):
+        return cfg
+    return build_from_cfg(cfg, BBOX_ASSIGNERS, default_args)
+
+
+def build_sampler(cfg, **default_args):
+    if isinstance(cfg, BaseSampler):
+        return cfg
+    return build_from_cfg(cfg, BBOX_SAMPLERS, default_args)
+
+
+# TODO remove this function in anchor_target in the future
+def assign_and_sample(bboxes, gt_bboxes, gt_bboxes_ignore, gt_labels, cfg):
+    bbox_assigner = build_assigner(cfg.assigner)
+    bbox_sampler = build_sampler(cfg.sampler)
+    assign_result = bbox_assigner.assign(bboxes, gt_bboxes, gt_bboxes_ignore,
+                                         gt_labels)
+    sampling_result = bbox_sampler.sample(assign_result, bboxes, gt_bboxes,
+                                          gt_labels)
+    return assign_result, sampling_result

mmdet/core/bbox/registry.py

@@ -0,0 +1,4 @@
+from mmdet.utils import Registry
+
+BBOX_ASSIGNERS = Registry('bbox_assigner')
+BBOX_SAMPLERS = Registry('bbox_sampler')

mmdet/core/bbox/samplers/combined_sampler.py

@@ -1,7 +1,9 @@
-from ..assign_sampling import build_sampler
+from ..builder import build_sampler
+from ..registry import BBOX_SAMPLERS
 from .base_sampler import BaseSampler
 
 
+@BBOX_SAMPLERS.register_module
 class CombinedSampler(BaseSampler):
 
     def __init__(self, pos_sampler, neg_sampler, **kwargs):

mmdet/core/bbox/samplers/instance_balanced_pos_sampler.py

@@ -1,9 +1,11 @@
 import numpy as np
 import torch
 
+from ..registry import BBOX_SAMPLERS
 from .random_sampler import RandomSampler
 
 
+@BBOX_SAMPLERS.register_module
 class InstanceBalancedPosSampler(RandomSampler):
 
     def _sample_pos(self, assign_result, num_expected, **kwargs):

mmdet/core/bbox/samplers/iou_balanced_neg_sampler.py

@@ -1,9 +1,11 @@
 import numpy as np
 import torch
 
+from ..registry import BBOX_SAMPLERS
 from .random_sampler import RandomSampler
 
 
+@BBOX_SAMPLERS.register_module
 class IoUBalancedNegSampler(RandomSampler):
     """IoU Balanced Sampling
 

mmdet/core/bbox/samplers/ohem_sampler.py

@@ -1,9 +1,11 @@
 import torch
 
+from ..registry import BBOX_SAMPLERS
 from ..transforms import bbox2roi
 from .base_sampler import BaseSampler
 
 
+@BBOX_SAMPLERS.register_module
 class OHEMSampler(BaseSampler):
     """
     Online Hard Example Mining Sampler described in [1]_.

mmdet/core/bbox/samplers/pseudo_sampler.py

@@ -1,9 +1,11 @@
 import torch
 
+from ..registry import BBOX_SAMPLERS
 from .base_sampler import BaseSampler
 from .sampling_result import SamplingResult
 
 
+@BBOX_SAMPLERS.register_module
 class PseudoSampler(BaseSampler):
 
     def __init__(self, **kwargs):

mmdet/core/bbox/samplers/random_sampler.py

@@ -1,8 +1,10 @@
 import torch
 
+from ..registry import BBOX_SAMPLERS
 from .base_sampler import BaseSampler
 
 
+@BBOX_SAMPLERS.register_module
 class RandomSampler(BaseSampler):
 
     def __init__(self,

tests/test_assigner.py

@@ -10,8 +10,8 @@
 """
 import torch
 
-from mmdet.core import MaxIoUAssigner
-from mmdet.core.bbox.assigners import ApproxMaxIoUAssigner, PointAssigner
+from mmdet.core.bbox.assigners import (ApproxMaxIoUAssigner, MaxIoUAssigner,
+                                       PointAssigner)
 
 
 def test_max_iou_assigner():

tests/test_sampler.py

@@ -1,6 +1,6 @@
 import torch
 
-from mmdet.core import MaxIoUAssigner
+from mmdet.core.bbox.assigners import MaxIoUAssigner
 from mmdet.core.bbox.samplers import OHEMSampler, RandomSampler