sparse_rcnn.py 1.2 KB

12345678910111213141516171819202122232425262728293031
  1. # Copyright (c) OpenMMLab. All rights reserved.
  2. from mmdet.registry import MODELS
  3. from mmdet.utils import ConfigType, OptConfigType, OptMultiConfig
  4. from .two_stage import TwoStageDetector
  5. @MODELS.register_module()
  6. class SparseRCNN(TwoStageDetector):
  7. r"""Implementation of `Sparse R-CNN: End-to-End Object Detection with
  8. Learnable Proposals <https://arxiv.org/abs/2011.12450>`_"""
  9. def __init__(self,
  10. backbone: ConfigType,
  11. neck: OptConfigType = None,
  12. rpn_head: OptConfigType = None,
  13. roi_head: OptConfigType = None,
  14. train_cfg: OptConfigType = None,
  15. test_cfg: OptConfigType = None,
  16. data_preprocessor: OptConfigType = None,
  17. init_cfg: OptMultiConfig = None) -> None:
  18. super().__init__(
  19. backbone=backbone,
  20. neck=neck,
  21. rpn_head=rpn_head,
  22. roi_head=roi_head,
  23. train_cfg=train_cfg,
  24. test_cfg=test_cfg,
  25. data_preprocessor=data_preprocessor,
  26. init_cfg=init_cfg)
  27. assert self.with_rpn, 'Sparse R-CNN and QueryInst ' \
  28. 'do not support external proposals'