Skip to content
Snippets Groups Projects
Commit 763decc4 authored by pangjm's avatar pangjm
Browse files

add fast rcnn api & fix minor bugs

parent b07edd6b
No related branches found
No related tags found
No related merge requests found
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']
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)
......@@ -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)
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment