From 5178d9f83244816dbce2c8bf17f1d3b7043237fc Mon Sep 17 00:00:00 2001
From: wangguanzhong <jerrywgz@126.com>
Date: Wed, 21 Oct 2020 18:02:30 +0800
Subject: [PATCH] support multiclass nms for multi-batch, test=develop (#28164)

---
 .../operators/detection/multiclass_nms_op.cc  | 10 +++-
 .../tests/unittests/test_multiclass_nms_op.py | 52 +++++++++++++++++++
 2 files changed, 61 insertions(+), 1 deletion(-)

diff --git a/paddle/fluid/operators/detection/multiclass_nms_op.cc b/paddle/fluid/operators/detection/multiclass_nms_op.cc
index 4b38779c136..0e835a62839 100644
--- a/paddle/fluid/operators/detection/multiclass_nms_op.cc
+++ b/paddle/fluid/operators/detection/multiclass_nms_op.cc
@@ -290,6 +290,7 @@ class MultiClassNMSKernel : public framework::OpKernel<T> {
       } else {
         sdata = scores_data + label * predict_dim;
       }
+
       for (size_t j = 0; j < indices.size(); ++j) {
         int idx = indices[j];
         odata[count * out_dim] = label;  // label
@@ -333,6 +334,7 @@ class MultiClassNMSKernel : public framework::OpKernel<T> {
     Tensor boxes_slice, scores_slice;
     int n = score_size == 3 ? batch_size : boxes->lod().back().size() - 1;
     for (int i = 0; i < n; ++i) {
+      std::map<int, std::vector<int>> indices;
       if (score_size == 3) {
         scores_slice = scores->Slice(i, i + 1);
         scores_slice.Resize({score_dims[1], score_dims[2]});
@@ -340,10 +342,14 @@ class MultiClassNMSKernel : public framework::OpKernel<T> {
         boxes_slice.Resize({score_dims[2], box_dim});
       } else {
         auto boxes_lod = boxes->lod().back();
+        if (boxes_lod[i] == boxes_lod[i + 1]) {
+          all_indices.push_back(indices);
+          batch_starts.push_back(batch_starts.back());
+          continue;
+        }
         scores_slice = scores->Slice(boxes_lod[i], boxes_lod[i + 1]);
         boxes_slice = boxes->Slice(boxes_lod[i], boxes_lod[i + 1]);
       }
-      std::map<int, std::vector<int>> indices;
       MultiClassNMS(ctx, scores_slice, boxes_slice, score_size, &indices,
                     &num_nmsed_out);
       all_indices.push_back(indices);
@@ -375,12 +381,14 @@ class MultiClassNMSKernel : public framework::OpKernel<T> {
           }
         } else {
           auto boxes_lod = boxes->lod().back();
+          if (boxes_lod[i] == boxes_lod[i + 1]) continue;
           scores_slice = scores->Slice(boxes_lod[i], boxes_lod[i + 1]);
           boxes_slice = boxes->Slice(boxes_lod[i], boxes_lod[i + 1]);
           if (return_index) {
             offset = boxes_lod[i] * score_dims[1];
           }
         }
+
         int64_t s = batch_starts[i];
         int64_t e = batch_starts[i + 1];
         if (e > s) {
diff --git a/python/paddle/fluid/tests/unittests/test_multiclass_nms_op.py b/python/paddle/fluid/tests/unittests/test_multiclass_nms_op.py
index ab58d4bc88e..34c19b88bcd 100644
--- a/python/paddle/fluid/tests/unittests/test_multiclass_nms_op.py
+++ b/python/paddle/fluid/tests/unittests/test_multiclass_nms_op.py
@@ -17,6 +17,7 @@ import unittest
 import numpy as np
 import copy
 from op_test import OpTest
+import paddle
 import paddle.fluid as fluid
 from paddle.fluid import Program, program_guard
 
@@ -171,6 +172,9 @@ def lod_multiclass_nms(boxes, scores, background, score_threshold,
     lod = []
     head = 0
     for n in range(len(box_lod[0])):
+        if box_lod[0][n] == 0:
+            lod.append(0)
+            continue
         box = boxes[head:head + box_lod[0][n]]
         score = scores[head:head + box_lod[0][n]]
         offset = head
@@ -357,6 +361,53 @@ class TestMulticlassNMSLoDInput(OpTest):
         self.check_output()
 
 
+class TestMulticlassNMSNoBox(TestMulticlassNMSLoDInput):
+    def setUp(self):
+        self.set_argument()
+        M = 1200
+        C = 21
+        BOX_SIZE = 4
+        box_lod = [[0, 1200, 0]]
+        background = 0
+        nms_threshold = 0.3
+        nms_top_k = 400
+        keep_top_k = 200
+        score_threshold = self.score_threshold
+        normalized = False
+
+        scores = np.random.random((M, C)).astype('float32')
+
+        scores = np.apply_along_axis(softmax, 1, scores)
+
+        boxes = np.random.random((M, C, BOX_SIZE)).astype('float32')
+        boxes[:, :, 0] = boxes[:, :, 0] * 10
+        boxes[:, :, 1] = boxes[:, :, 1] * 10
+        boxes[:, :, 2] = boxes[:, :, 2] * 10 + 10
+        boxes[:, :, 3] = boxes[:, :, 3] * 10 + 10
+
+        det_outs, lod = lod_multiclass_nms(
+            boxes, scores, background, score_threshold, nms_threshold,
+            nms_top_k, keep_top_k, box_lod, normalized)
+        det_outs = np.array(det_outs).astype('float32')
+        nmsed_outs = det_outs[:, :-1].astype('float32') if len(
+            det_outs) else det_outs
+        self.op_type = 'multiclass_nms'
+        self.inputs = {
+            'BBoxes': (boxes, box_lod),
+            'Scores': (scores, box_lod),
+        }
+        self.outputs = {'Out': (nmsed_outs, [lod])}
+        self.attrs = {
+            'background_label': 0,
+            'nms_threshold': nms_threshold,
+            'nms_top_k': nms_top_k,
+            'keep_top_k': keep_top_k,
+            'score_threshold': score_threshold,
+            'nms_eta': 1.0,
+            'normalized': normalized,
+        }
+
+
 class TestIOU(unittest.TestCase):
     def test_iou(self):
         box1 = np.array([4.0, 3.0, 7.0, 5.0]).astype('float32')
@@ -521,4 +572,5 @@ class TestMulticlassNMSError(unittest.TestCase):
 
 
 if __name__ == '__main__':
+    paddle.enable_static()
     unittest.main()
-- 
GitLab