diff --git a/python/paddle/fluid/layers/detection.py b/python/paddle/fluid/layers/detection.py index 18f057623a393509370dbf3b6501f17735f0adc6..cd179f40f47d881267ac45302f8e8ef79a49784e 100644 --- a/python/paddle/fluid/layers/detection.py +++ b/python/paddle/fluid/layers/detection.py @@ -2742,7 +2742,8 @@ def generate_proposals(scores, nms_thresh=0.5, min_size=0.1, eta=1.0, - name=None): + name=None, + return_rois_num=False): """ **Generate proposal Faster-RCNN** @@ -2789,7 +2790,10 @@ def generate_proposals(scores, width < min_size. The data type must be float32. `0.1` by default. eta(float): Apply in adaptive NMS, if adaptive `threshold > 0.5`, `adaptive_threshold = adaptive_threshold * eta` in each iteration. - + return_rois_num(bool): When setting True, it will return a 1D Tensor with shape [N, ] that includes Rois's + num of each image in one batch. The N is the image's num. For example, the tensor has values [4,5] that represents + the first image has 4 Rois, the second image has 5 Rois. It only used in rcnn model. + 'False' by default. Returns: tuple: A tuple with format ``(rpn_rois, rpn_roi_probs)``. @@ -2843,7 +2847,10 @@ def generate_proposals(scores, rpn_roi_probs.stop_gradient = True rpn_rois_lod.stop_gradient = True - return rpn_rois, rpn_roi_probs, rpn_rois_lod + if return_rois_num: + return rpn_rois, rpn_roi_probs, rpn_rois_lod + else: + return rpn_rois, rpn_roi_probs def box_clip(input, im_info, name=None): diff --git a/python/paddle/fluid/tests/test_detection.py b/python/paddle/fluid/tests/test_detection.py index 0a220fe18f42b27021be4091293c8baa8f038e53..cb798c8ed595d13dd8ff5e33323d6e796aaac6f9 100644 --- a/python/paddle/fluid/tests/test_detection.py +++ b/python/paddle/fluid/tests/test_detection.py @@ -480,7 +480,7 @@ class TestGenerateProposals(unittest.TestCase): name='bbox_deltas', shape=[num_anchors * 4, 8, 8], dtype='float32') - rpn_rois, rpn_roi_probs, _ = fluid.layers.generate_proposals( + rpn_rois, rpn_roi_probs = fluid.layers.generate_proposals( name='generate_proposals', scores=scores, bbox_deltas=bbox_deltas, diff --git a/python/paddle/fluid/tests/unittests/test_generate_proposals_op.py b/python/paddle/fluid/tests/unittests/test_generate_proposals_op.py index 5d4b453fa64205cb13b88ccfc5bf622348bbd9fa..ce561cd317c48228c4877a2b65b67fe049a0d84a 100644 --- a/python/paddle/fluid/tests/unittests/test_generate_proposals_op.py +++ b/python/paddle/fluid/tests/unittests/test_generate_proposals_op.py @@ -282,8 +282,6 @@ class TestGenerateProposalsOp(OpTest): self.outputs = { 'RpnRois': (self.rpn_rois[0], [self.lod]), 'RpnRoiProbs': (self.rpn_roi_probs[0], [self.lod]), - 'RpnRoisLod': (np.asarray( - self.lod, dtype=np.int32)) } def test_check_output(self): @@ -328,5 +326,35 @@ class TestGenerateProposalsOp(OpTest): self.nms_thresh, self.min_size, self.eta) +class TestGenerateProposalsOutLodOp(TestGenerateProposalsOp): + def set_data(self): + self.init_test_params() + self.init_test_input() + self.init_test_output() + self.inputs = { + 'Scores': self.scores, + 'BboxDeltas': self.bbox_deltas, + 'ImInfo': self.im_info.astype(np.float32), + 'Anchors': self.anchors, + 'Variances': self.variances + } + + self.attrs = { + 'pre_nms_topN': self.pre_nms_topN, + 'post_nms_topN': self.post_nms_topN, + 'nms_thresh': self.nms_thresh, + 'min_size': self.min_size, + 'eta': self.eta, + 'return_rois_num': True + } + + self.outputs = { + 'RpnRois': (self.rpn_rois[0], [self.lod]), + 'RpnRoiProbs': (self.rpn_roi_probs[0], [self.lod]), + 'RpnRoisLod': (np.asarray( + self.lod, dtype=np.int32)) + } + + if __name__ == '__main__': unittest.main()