From b62a17bbae254c0b96169cab0129dd942ff19083 Mon Sep 17 00:00:00 2001
From: jerrywgz <jerrywgz@126.com>
Date: Fri, 18 Jan 2019 10:01:47 +0000
Subject: [PATCH] add nms api

---
 .../operators/detection/multiclass_nms_op.cc  |  8 ++---
 python/paddle/fluid/layers/detection.py       | 35 +++++++++++++++++++
 python/paddle/fluid/tests/test_detection.py   | 11 ++++++
 3 files changed, 50 insertions(+), 4 deletions(-)

diff --git a/paddle/fluid/operators/detection/multiclass_nms_op.cc b/paddle/fluid/operators/detection/multiclass_nms_op.cc
index 14ce9937d..c61e3e133 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<int>(
         "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<bool>("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 8aed97dc5..e8ce0c1d9 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 d99eaa063..7736cfc2f 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()
-- 
GitLab