diff --git a/paddle/fluid/operators/detection/multiclass_nms_op.cc b/paddle/fluid/operators/detection/multiclass_nms_op.cc index 14ce9937dc69348c80d58abf07db5db7fc31cdbe..c61e3e1338dc39939480db9d267455e80fd0d216 100644 --- a/paddle/fluid/operators/detection/multiclass_nms_op.cc +++ b/paddle/fluid/operators/detection/multiclass_nms_op.cc @@ -458,7 +458,8 @@ class MultiClassNMSOpMaker : public framework::OpProtoAndCheckerMaker { "predicted locations of M bounding bboxes, N is the batch size. " "Each bounding box has four coordinate values and the layout is " "[xmin, ymin, xmax, ymax], when box size equals to 4." - "2. (LoDTensor) A 3-D Tensor with shape [N, M, 4]"); + "2. (LoDTensor) A 3-D Tensor with shape [N, M, 4]" + "N is the number of boxes, M is the class number"); AddInput("Scores", "Two types of scores are supported:" "1. (Tensor) A 3-D Tensor with shape [N, C, M] represents the " @@ -467,8 +468,7 @@ class MultiClassNMSOpMaker : public framework::OpProtoAndCheckerMaker { "there are total M scores which corresponding M bounding boxes. " " Please note, M is equal to the 1st dimension of BBoxes. " "2. (LoDTensor) A 2-D LoDTensor with shape" - "[N, num_class]. N is the number of bbox and" - "M represents the scores of bboxes in each class."); + "[N, num_class]. N is the number of bbox"); AddAttr( "background_label", "(int, defalut: 0) " @@ -497,7 +497,7 @@ class MultiClassNMSOpMaker : public framework::OpProtoAndCheckerMaker { "Number of total bboxes to be kept per image after NMS " "step. -1 means keeping all bboxes after NMS step."); AddAttr("normalized", - "(bool, default false) " + "(bool, default true) " "Whether detections are normalized.") .SetDefault(true); AddOutput("Out", diff --git a/python/paddle/fluid/layers/detection.py b/python/paddle/fluid/layers/detection.py index 8aed97dc59b100d4e37832e0a148d73662742ba0..e8ce0c1d90db430d6fe320fd6493a2179c5584f9 100644 --- a/python/paddle/fluid/layers/detection.py +++ b/python/paddle/fluid/layers/detection.py @@ -48,6 +48,7 @@ __all__ = [ 'box_coder', 'polygon_box_transform', 'yolov3_loss', + 'multiclass_nms', ] @@ -1810,3 +1811,37 @@ def generate_proposals(scores, rpn_roi_probs.stop_gradient = True return rpn_rois, rpn_roi_probs + + +def multiclass_nms(bboxes, + scores, + score_threshold, + nms_top_k, + nms_threshold, + keep_top_k, + normalized=True, + nms_eta=1., + background_label=0): + """ + """ + helper = LayerHelper('multiclass_nms', **locals()) + + output = helper.create_variable_for_type_inference(dtype=bboxes.dtype) + helper.append_op( + type="multiclass_nms", + inputs={'BBoxes': bboxes, + 'Scores': scores}, + attrs={ + 'background_label': background_label, + 'score_threshold': score_threshold, + 'nms_top_k': nms_top_k, + 'nms_threshold': nms_threshold, + 'nms_eta': nms_eta, + 'keep_top_k': keep_top_k, + 'nms_eta': nms_eta, + 'normalized': normalized + }, + outputs={'Out': output}) + output.stop_gradient = True + + return output diff --git a/python/paddle/fluid/tests/test_detection.py b/python/paddle/fluid/tests/test_detection.py index d99eaa0634f93dcd16dd80ae172f11e8090a2623..7736cfc2fb7ecd87f154109c30abcf010e196bfe 100644 --- a/python/paddle/fluid/tests/test_detection.py +++ b/python/paddle/fluid/tests/test_detection.py @@ -401,5 +401,16 @@ class TestYoloDetection(unittest.TestCase): self.assertIsNotNone(loss) +class TestMulticlassNMS(unittest.TestCase): + def test_multiclass_nms(self): + program = Program() + with program_guard(program): + bboxes = layers.data( + name='bboxes', shape=[-1, 10, 4], dtype='float32') + scores = layers.data(name='scores', shape=[-1, 10], dtype='float32') + output = layers.multiclass_nms(bboxes, scores, 0.3, 400, 0.7, 200) + self.assertIsNotNone(output) + + if __name__ == '__main__': unittest.main()