diff --git a/mmdet/models/detectors/__init__.py b/mmdet/models/detectors/__init__.py index b8914c1e5d3c834a1373b2a2e8360183a41de4da..29a64dd09b69f9d617053fea5578c1f9b1170e56 100644 --- a/mmdet/models/detectors/__init__.py +++ b/mmdet/models/detectors/__init__.py @@ -1,6 +1,7 @@ from .base import BaseDetector from .rpn import RPN +from .faster_rcnn import FastRCNN from .faster_rcnn import FasterRCNN from .mask_rcnn import MaskRCNN -__all__ = ['BaseDetector', 'RPN', 'FasterRCNN', 'MaskRCNN'] +__all__ = ['BaseDetector', 'RPN', 'FastRCNN', 'FasterRCNN', 'MaskRCNN'] diff --git a/mmdet/models/detectors/fast_rcnn.py b/mmdet/models/detectors/fast_rcnn.py new file mode 100644 index 0000000000000000000000000000000000000000..0dbf17a9ab9bc7ff692026f3c9c35364384c0df2 --- /dev/null +++ b/mmdet/models/detectors/fast_rcnn.py @@ -0,0 +1,25 @@ +from .two_stage import TwoStageDetector + + +class FastRCNN(TwoStageDetector): + + def __init__(self, + backbone, + neck, + bbox_roi_extractor, + bbox_head, + train_cfg, + test_cfg, + mask_roi_extractor=None, + mask_head=None, + pretrained=None): + super(FastRCNN, self).__init__( + backbone=backbone, + neck=neck, + bbox_roi_extractor=bbox_roi_extractor, + bbox_head=bbox_head, + train_cfg=train_cfg, + test_cfg=test_cfg, + mask_roi_extractor=mask_roi_extractor, + mask_head=mask_head, + pretrained=pretrained) diff --git a/mmdet/models/detectors/two_stage.py b/mmdet/models/detectors/two_stage.py index 8573d83215f120ba392a2f6b45cb9b6b93ca0519..3cd68388790b164c70b8d88708d0e1a170529ab4 100644 --- a/mmdet/models/detectors/two_stage.py +++ b/mmdet/models/detectors/two_stage.py @@ -146,7 +146,8 @@ class TwoStageDetector(BaseDetector, RPNTestMixin, BBoxTestMixin, x = self.extract_feat(img) proposal_list = self.simple_test_rpn( - x, img_meta, self.test_cfg.rpn) if proposals is None else proposals + x, img_meta, + self.test_cfg.rpn) if proposals is None else proposals[0] det_bboxes, det_labels = self.simple_test_bboxes( x, img_meta, proposal_list, self.test_cfg.rcnn, rescale=rescale)