multiclass_nms_op.cc 24.6 KB
Newer Older
1
/* Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
2

3 4 5
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
6

7
http://www.apache.org/licenses/LICENSE-2.0
8

9 10 11 12 13
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
limitations under the License. */

J
jerrywgz 已提交
14
#include <glog/logging.h>
Y
Yi Wang 已提交
15
#include "paddle/fluid/framework/op_registry.h"
16
#include "paddle/fluid/operators/detection/nms_util.h"
17 18 19 20 21 22 23

namespace paddle {
namespace operators {

using Tensor = framework::Tensor;
using LoDTensor = framework::LoDTensor;

24 25 26 27 28 29 30 31 32 33
inline std::vector<size_t> GetNmsLodFromRoisNum(const Tensor* rois_num) {
  std::vector<size_t> rois_lod;
  auto* rois_num_data = rois_num->data<int>();
  rois_lod.push_back(static_cast<size_t>(0));
  for (int i = 0; i < rois_num->numel(); ++i) {
    rois_lod.push_back(rois_lod.back() + static_cast<size_t>(rois_num_data[i]));
  }
  return rois_lod;
}

D
dangqingqing 已提交
34
class MultiClassNMSOp : public framework::OperatorWithKernel {
35 36 37 38
 public:
  using framework::OperatorWithKernel::OperatorWithKernel;

  void InferShape(framework::InferShapeContext* ctx) const override {
X
xiaoting 已提交
39 40 41
    OP_INOUT_CHECK(ctx->HasInput("BBoxes"), "Input", "BBoxes", "MultiClassNMS");
    OP_INOUT_CHECK(ctx->HasInput("Scores"), "Input", "Scores", "MultiClassNMS");
    OP_INOUT_CHECK(ctx->HasOutput("Out"), "Output", "Out", "MultiClassNMS");
D
dangqingqing 已提交
42
    auto box_dims = ctx->GetInputDim("BBoxes");
43
    auto score_dims = ctx->GetInputDim("Scores");
J
jerrywgz 已提交
44
    auto score_size = score_dims.size();
45

46
    if (ctx->IsRuntime()) {
47 48 49 50 51
      PADDLE_ENFORCE_EQ(score_size == 2 || score_size == 3, true,
                        platform::errors::InvalidArgument(
                            "The rank of Input(Scores) must be 2 or 3"
                            ". But received rank = %d",
                            score_size));
52
      PADDLE_ENFORCE_EQ(box_dims.size(), 3,
X
xiaoting 已提交
53 54
                        platform::errors::InvalidArgument(
                            "The rank of Input(BBoxes) must be 3"
55
                            ". But received rank = %d",
X
xiaoting 已提交
56
                            box_dims.size()));
J
jerrywgz 已提交
57
      if (score_size == 3) {
58 59 60 61 62 63 64 65 66 67 68 69
        PADDLE_ENFORCE_EQ(
            box_dims[2] == 4 || box_dims[2] == 8 || box_dims[2] == 16 ||
                box_dims[2] == 24 || box_dims[2] == 32,
            true, platform::errors::InvalidArgument(
                      "The last dimension of Input"
                      "(BBoxes) must be 4 or 8, "
                      "represents the layout of coordinate "
                      "[xmin, ymin, xmax, ymax] or "
                      "4 points: [x1, y1, x2, y2, x3, y3, x4, y4] or "
                      "8 points: [xi, yi] i= 1,2,...,8 or "
                      "12 points: [xi, yi] i= 1,2,...,12 or "
                      "16 points: [xi, yi] i= 1,2,...,16"));
J
jerrywgz 已提交
70 71
        PADDLE_ENFORCE_EQ(
            box_dims[1], score_dims[2],
X
xiaoting 已提交
72 73 74 75 76 77
            platform::errors::InvalidArgument(
                "The 2nd dimension of Input(BBoxes) must be equal to "
                "last dimension of Input(Scores), which represents the "
                "predicted bboxes."
                "But received box_dims[1](%s) != socre_dims[2](%s)",
                box_dims[1], score_dims[2]));
J
jerrywgz 已提交
78
      } else {
X
xiaoting 已提交
79 80
        PADDLE_ENFORCE_EQ(box_dims[2], 4,
                          platform::errors::InvalidArgument(
81 82
                              "The last dimension of Input"
                              "(BBoxes) must be 4. But received dimension = %d",
X
xiaoting 已提交
83
                              box_dims[2]));
84 85 86 87 88 89 90
        PADDLE_ENFORCE_EQ(
            box_dims[1], score_dims[1],
            platform::errors::InvalidArgument(
                "The 2nd dimension of Input"
                "(BBoxes) must be equal to the 2nd dimension of Input(Scores). "
                "But received box dimension = %d, score dimension = %d",
                box_dims[1], score_dims[1]));
J
jerrywgz 已提交
91
      }
92
    }
93 94
    // Here the box_dims[0] is not the real dimension of output.
    // It will be rewritten in the computing kernel.
J
jerrywgz 已提交
95
    if (score_size == 3) {
96
      ctx->SetOutputDim("Out", {-1, box_dims[2] + 2});
J
jerrywgz 已提交
97 98 99
    } else {
      ctx->SetOutputDim("Out", {-1, box_dims[2] + 2});
    }
100 101 102
    if (!ctx->IsRuntime()) {
      ctx->SetLoDLevel("Out", std::max(ctx->GetLoDLevel("BBoxes"), 1));
    }
103
  }
D
dangqingqing 已提交
104 105 106 107 108

 protected:
  framework::OpKernelType GetExpectedKernelType(
      const framework::ExecutionContext& ctx) const override {
    return framework::OpKernelType(
109
        OperatorWithKernel::IndicateVarDataType(ctx, "Scores"),
110
        platform::CPUPlace());
D
dangqingqing 已提交
111
  }
112 113
};

114 115 116 117 118 119 120 121 122
template <class T>
void SliceOneClass(const platform::DeviceContext& ctx,
                   const framework::Tensor& items, const int class_id,
                   framework::Tensor* one_class_item) {
  T* item_data = one_class_item->mutable_data<T>(ctx.GetPlace());
  const T* items_data = items.data<T>();
  const int64_t num_item = items.dims()[0];
  const int class_num = items.dims()[1];
  if (items.dims().size() == 3) {
J
jerrywgz 已提交
123 124 125 126 127 128 129 130 131 132
    int item_size = items.dims()[2];
    for (int i = 0; i < num_item; ++i) {
      std::memcpy(item_data + i * item_size,
                  items_data + i * class_num * item_size + class_id * item_size,
                  sizeof(T) * item_size);
    }
  } else {
    for (int i = 0; i < num_item; ++i) {
      item_data[i] = items_data[i * class_num + class_id];
    }
133 134 135
  }
}

136
template <typename T>
D
dangqingqing 已提交
137
class MultiClassNMSKernel : public framework::OpKernel<T> {
138 139 140
 public:
  void NMSFast(const Tensor& bbox, const Tensor& scores,
               const T score_threshold, const T nms_threshold, const T eta,
J
jerrywgz 已提交
141 142
               const int64_t top_k, std::vector<int>* selected_indices,
               const bool normalized) const {
143 144 145
    // The total boxes for each instance.
    int64_t num_boxes = bbox.dims()[0];
    // 4: [xmin ymin xmax ymax]
Y
Yipeng 已提交
146 147
    // 8: [x1 y1 x2 y2 x3 y3 x4 y4]
    // 16, 24, or 32: [x1 y1 x2 y2 ...  xn yn], n = 8, 12 or 16
148 149 150 151 152 153 154 155 156 157 158 159 160 161
    int64_t box_size = bbox.dims()[1];

    std::vector<T> scores_data(num_boxes);
    std::copy_n(scores.data<T>(), num_boxes, scores_data.begin());
    std::vector<std::pair<T, int>> sorted_indices;
    GetMaxScoreIndex(scores_data, score_threshold, top_k, &sorted_indices);

    selected_indices->clear();
    T adaptive_threshold = nms_threshold;
    const T* bbox_data = bbox.data<T>();

    while (sorted_indices.size() != 0) {
      const int idx = sorted_indices.front().second;
      bool keep = true;
162
      for (size_t k = 0; k < selected_indices->size(); ++k) {
163 164
        if (keep) {
          const int kept_idx = (*selected_indices)[k];
Y
Yipeng 已提交
165 166 167
          T overlap = T(0.);
          // 4: [xmin ymin xmax ymax]
          if (box_size == 4) {
J
jerrywgz 已提交
168 169 170
            overlap =
                JaccardOverlap<T>(bbox_data + idx * box_size,
                                  bbox_data + kept_idx * box_size, normalized);
Y
Yipeng 已提交
171 172 173 174
          }
          // 8: [x1 y1 x2 y2 x3 y3 x4 y4] or 16, 24, 32
          if (box_size == 8 || box_size == 16 || box_size == 24 ||
              box_size == 32) {
J
jerrywgz 已提交
175 176 177
            overlap = PolyIoU<T>(bbox_data + idx * box_size,
                                 bbox_data + kept_idx * box_size, box_size,
                                 normalized);
Y
Yipeng 已提交
178
          }
179 180 181 182 183 184 185 186 187 188 189 190 191 192 193
          keep = overlap <= adaptive_threshold;
        } else {
          break;
        }
      }
      if (keep) {
        selected_indices->push_back(idx);
      }
      sorted_indices.erase(sorted_indices.begin());
      if (keep && eta < 1 && adaptive_threshold > 0.5) {
        adaptive_threshold *= eta;
      }
    }
  }

D
dangqingqing 已提交
194
  void MultiClassNMS(const framework::ExecutionContext& ctx,
195
                     const Tensor& scores, const Tensor& bboxes,
J
jerrywgz 已提交
196
                     const int scores_size,
197 198
                     std::map<int, std::vector<int>>* indices,
                     int* num_nmsed_out) const {
D
dangqingqing 已提交
199 200 201
    int64_t background_label = ctx.Attr<int>("background_label");
    int64_t nms_top_k = ctx.Attr<int>("nms_top_k");
    int64_t keep_top_k = ctx.Attr<int>("keep_top_k");
J
jerrywgz 已提交
202
    bool normalized = ctx.Attr<bool>("normalized");
203 204
    T nms_threshold = static_cast<T>(ctx.Attr<float>("nms_threshold"));
    T nms_eta = static_cast<T>(ctx.Attr<float>("nms_eta"));
D
dangqingqing 已提交
205
    T score_threshold = static_cast<T>(ctx.Attr<float>("score_threshold"));
J
jerrywgz 已提交
206
    auto& dev_ctx = ctx.template device_context<platform::CPUDeviceContext>();
207 208

    int num_det = 0;
209 210 211 212 213 214 215 216 217 218 219 220 221

    int64_t class_num = scores_size == 3 ? scores.dims()[0] : scores.dims()[1];
    Tensor bbox_slice, score_slice;
    for (int64_t c = 0; c < class_num; ++c) {
      if (c == background_label) continue;
      if (scores_size == 3) {
        score_slice = scores.Slice(c, c + 1);
        bbox_slice = bboxes;
      } else {
        score_slice.Resize({scores.dims()[0], 1});
        bbox_slice.Resize({scores.dims()[0], 4});
        SliceOneClass<T>(dev_ctx, scores, c, &score_slice);
        SliceOneClass<T>(dev_ctx, bboxes, c, &bbox_slice);
J
jerrywgz 已提交
222
      }
223 224 225
      NMSFast(bbox_slice, score_slice, score_threshold, nms_threshold, nms_eta,
              nms_top_k, &((*indices)[c]), normalized);
      if (scores_size == 2) {
J
jerrywgz 已提交
226 227
        std::stable_sort((*indices)[c].begin(), (*indices)[c].end());
      }
228
      num_det += (*indices)[c].size();
229 230
    }

231
    *num_nmsed_out = num_det;
232 233
    const T* scores_data = scores.data<T>();
    if (keep_top_k > -1 && num_det > keep_top_k) {
J
jerrywgz 已提交
234
      const T* sdata;
235
      std::vector<std::pair<float, std::pair<int, int>>> score_index_pairs;
236
      for (const auto& it : *indices) {
237
        int label = it.first;
J
jerrywgz 已提交
238
        if (scores_size == 3) {
239
          sdata = scores_data + label * scores.dims()[1];
J
jerrywgz 已提交
240
        } else {
241 242 243
          score_slice.Resize({scores.dims()[0], 1});
          SliceOneClass<T>(dev_ctx, scores, label, &score_slice);
          sdata = score_slice.data<T>();
J
jerrywgz 已提交
244
        }
245
        const std::vector<int>& label_indices = it.second;
246
        for (size_t j = 0; j < label_indices.size(); ++j) {
247 248 249 250 251 252
          int idx = label_indices[j];
          score_index_pairs.push_back(
              std::make_pair(sdata[idx], std::make_pair(label, idx)));
        }
      }
      // Keep top k results per image.
253 254
      std::stable_sort(score_index_pairs.begin(), score_index_pairs.end(),
                       SortScorePairDescend<std::pair<int, int>>);
255 256 257 258
      score_index_pairs.resize(keep_top_k);

      // Store the new indices.
      std::map<int, std::vector<int>> new_indices;
259
      for (size_t j = 0; j < score_index_pairs.size(); ++j) {
260 261 262 263
        int label = score_index_pairs[j].second.first;
        int idx = score_index_pairs[j].second.second;
        new_indices[label].push_back(idx);
      }
J
jerrywgz 已提交
264 265 266 267 268 269 270
      if (scores_size == 2) {
        for (const auto& it : new_indices) {
          int label = it.first;
          std::stable_sort(new_indices[label].begin(),
                           new_indices[label].end());
        }
      }
271 272
      new_indices.swap(*indices);
      *num_nmsed_out = keep_top_k;
273 274 275
    }
  }

J
jerrywgz 已提交
276 277
  void MultiClassOutput(const platform::DeviceContext& ctx,
                        const Tensor& scores, const Tensor& bboxes,
278
                        const std::map<int, std::vector<int>>& selected_indices,
279 280
                        const int scores_size, Tensor* outs,
                        int* oindices = nullptr, const int offset = 0) const {
J
jerrywgz 已提交
281
    int64_t class_num = scores.dims()[1];
Y
Yipeng 已提交
282 283
    int64_t predict_dim = scores.dims()[1];
    int64_t box_size = bboxes.dims()[1];
J
jerrywgz 已提交
284 285 286 287
    if (scores_size == 2) {
      box_size = bboxes.dims()[2];
    }
    int64_t out_dim = box_size + 2;
288 289 290
    auto* scores_data = scores.data<T>();
    auto* bboxes_data = bboxes.data<T>();
    auto* odata = outs->data<T>();
J
jerrywgz 已提交
291 292 293
    const T* sdata;
    Tensor bbox;
    bbox.Resize({scores.dims()[0], box_size});
294 295 296
    int count = 0;
    for (const auto& it : selected_indices) {
      int label = it.first;
D
dangqingqing 已提交
297
      const std::vector<int>& indices = it.second;
J
jerrywgz 已提交
298 299 300 301 302
      if (scores_size == 2) {
        SliceOneClass<T>(ctx, bboxes, label, &bbox);
      } else {
        sdata = scores_data + label * predict_dim;
      }
303

304
      for (size_t j = 0; j < indices.size(); ++j) {
305
        int idx = indices[j];
J
jerrywgz 已提交
306 307 308 309 310
        odata[count * out_dim] = label;  // label
        const T* bdata;
        if (scores_size == 3) {
          bdata = bboxes_data + idx * box_size;
          odata[count * out_dim + 1] = sdata[idx];  // score
311 312 313
          if (oindices != nullptr) {
            oindices[count] = offset + idx;
          }
J
jerrywgz 已提交
314 315 316
        } else {
          bdata = bbox.data<T>() + idx * box_size;
          odata[count * out_dim + 1] = *(scores_data + idx * class_num + label);
317 318 319
          if (oindices != nullptr) {
            oindices[count] = offset + idx * class_num + label;
          }
J
jerrywgz 已提交
320
        }
Y
Yipeng 已提交
321 322
        // xmin, ymin, xmax, ymax or multi-points coordinates
        std::memcpy(odata + count * out_dim + 2, bdata, box_size * sizeof(T));
D
dangqingqing 已提交
323
        count++;
324 325 326 327 328
      }
    }
  }

  void Compute(const framework::ExecutionContext& ctx) const override {
J
jerrywgz 已提交
329 330
    auto* boxes = ctx.Input<LoDTensor>("BBoxes");
    auto* scores = ctx.Input<LoDTensor>("Scores");
331
    auto* outs = ctx.Output<LoDTensor>("Out");
332 333
    bool return_index = ctx.HasOutput("Index") ? true : false;
    auto index = ctx.Output<LoDTensor>("Index");
334 335
    bool has_roisnum = ctx.HasInput("RoisNum") ? true : false;
    auto rois_num = ctx.Input<Tensor>("RoisNum");
336
    auto score_dims = scores->dims();
337
    auto score_size = score_dims.size();
J
jerrywgz 已提交
338
    auto& dev_ctx = ctx.template device_context<platform::CPUDeviceContext>();
339 340 341

    std::vector<std::map<int, std::vector<int>>> all_indices;
    std::vector<size_t> batch_starts = {0};
J
jerrywgz 已提交
342 343 344 345
    int64_t batch_size = score_dims[0];
    int64_t box_dim = boxes->dims()[2];
    int64_t out_dim = box_dim + 2;
    int num_nmsed_out = 0;
346
    Tensor boxes_slice, scores_slice;
347 348 349 350 351 352
    int n = 0;
    if (has_roisnum) {
      n = score_size == 3 ? batch_size : rois_num->numel();
    } else {
      n = score_size == 3 ? batch_size : boxes->lod().back().size() - 1;
    }
353
    for (int i = 0; i < n; ++i) {
354
      std::map<int, std::vector<int>> indices;
355 356 357 358 359 360
      if (score_size == 3) {
        scores_slice = scores->Slice(i, i + 1);
        scores_slice.Resize({score_dims[1], score_dims[2]});
        boxes_slice = boxes->Slice(i, i + 1);
        boxes_slice.Resize({score_dims[2], box_dim});
      } else {
361 362 363 364 365 366
        std::vector<size_t> boxes_lod;
        if (has_roisnum) {
          boxes_lod = GetNmsLodFromRoisNum(rois_num);
        } else {
          boxes_lod = boxes->lod().back();
        }
367 368 369 370 371
        if (boxes_lod[i] == boxes_lod[i + 1]) {
          all_indices.push_back(indices);
          batch_starts.push_back(batch_starts.back());
          continue;
        }
372 373
        scores_slice = scores->Slice(boxes_lod[i], boxes_lod[i + 1]);
        boxes_slice = boxes->Slice(boxes_lod[i], boxes_lod[i + 1]);
J
jerrywgz 已提交
374
      }
375 376 377 378
      MultiClassNMS(ctx, scores_slice, boxes_slice, score_size, &indices,
                    &num_nmsed_out);
      all_indices.push_back(indices);
      batch_starts.push_back(batch_starts.back() + num_nmsed_out);
J
jerrywgz 已提交
379 380 381 382
    }

    int num_kept = batch_starts.back();
    if (num_kept == 0) {
383 384 385 386 387 388 389 390
      if (return_index) {
        outs->mutable_data<T>({0, out_dim}, ctx.GetPlace());
        index->mutable_data<int>({0, 1}, ctx.GetPlace());
      } else {
        T* od = outs->mutable_data<T>({1, 1}, ctx.GetPlace());
        od[0] = -1;
        batch_starts = {0, 1};
      }
J
jerrywgz 已提交
391 392
    } else {
      outs->mutable_data<T>({num_kept, out_dim}, ctx.GetPlace());
393 394
      int offset = 0;
      int* oindices = nullptr;
395 396 397 398 399 400
      for (int i = 0; i < n; ++i) {
        if (score_size == 3) {
          scores_slice = scores->Slice(i, i + 1);
          boxes_slice = boxes->Slice(i, i + 1);
          scores_slice.Resize({score_dims[1], score_dims[2]});
          boxes_slice.Resize({score_dims[2], box_dim});
401 402 403
          if (return_index) {
            offset = i * score_dims[2];
          }
404
        } else {
405 406 407 408 409 410
          std::vector<size_t> boxes_lod;
          if (has_roisnum) {
            boxes_lod = GetNmsLodFromRoisNum(rois_num);
          } else {
            boxes_lod = boxes->lod().back();
          }
411
          if (boxes_lod[i] == boxes_lod[i + 1]) continue;
412 413
          scores_slice = scores->Slice(boxes_lod[i], boxes_lod[i + 1]);
          boxes_slice = boxes->Slice(boxes_lod[i], boxes_lod[i + 1]);
414 415 416
          if (return_index) {
            offset = boxes_lod[i] * score_dims[1];
          }
J
jerrywgz 已提交
417
        }
418

419 420 421 422
        int64_t s = batch_starts[i];
        int64_t e = batch_starts[i + 1];
        if (e > s) {
          Tensor out = outs->Slice(s, e);
423 424 425 426 427
          if (return_index) {
            int* output_idx =
                index->mutable_data<int>({num_kept, 1}, ctx.GetPlace());
            oindices = output_idx + s;
          }
428
          MultiClassOutput(dev_ctx, scores_slice, boxes_slice, all_indices[i],
429
                           score_dims.size(), &out, oindices, offset);
430 431 432
        }
      }
    }
433 434 435 436 437 438 439 440 441
    if (ctx.HasOutput("NmsRoisNum")) {
      auto* nms_rois_num = ctx.Output<Tensor>("NmsRoisNum");
      nms_rois_num->mutable_data<int>({n}, ctx.GetPlace());
      int* num_data = nms_rois_num->data<int>();
      for (int i = 1; i <= n; i++) {
        num_data[i - 1] = batch_starts[i] - batch_starts[i - 1];
      }
      nms_rois_num->Resize({n});
    }
442 443 444

    framework::LoD lod;
    lod.emplace_back(batch_starts);
445 446 447
    if (return_index) {
      index->set_lod(lod);
    }
448 449 450 451
    outs->set_lod(lod);
  }
};

D
dangqingqing 已提交
452
class MultiClassNMSOpMaker : public framework::OpProtoAndCheckerMaker {
453
 public:
Y
Yu Yang 已提交
454
  void Make() override {
D
dangqingqing 已提交
455
    AddInput("BBoxes",
J
jerrywgz 已提交
456 457
             "Two types of bboxes are supported:"
             "1. (Tensor) A 3-D Tensor with shape "
Y
Yipeng 已提交
458
             "[N, M, 4 or 8 16 24 32] represents the "
459 460
             "predicted locations of M bounding bboxes, N is the batch size. "
             "Each bounding box has four coordinate values and the layout is "
J
jerrywgz 已提交
461
             "[xmin, ymin, xmax, ymax], when box size equals to 4."
462 463
             "2. (LoDTensor) A 3-D Tensor with shape [M, C, 4]"
             "M is the number of bounding boxes, C is the class number");
D
dangqingqing 已提交
464
    AddInput("Scores",
J
jerrywgz 已提交
465 466
             "Two types of scores are supported:"
             "1. (Tensor) A 3-D Tensor with shape [N, C, M] represents the "
D
dangqingqing 已提交
467 468 469
             "predicted confidence predictions. N is the batch size, C is the "
             "class number, M is number of bounding boxes. For each category "
             "there are total M scores which corresponding M bounding boxes. "
470 471 472 473
             " Please note, M is equal to the 2nd dimension of BBoxes. "
             "2. (LoDTensor) A 2-D LoDTensor with shape [M, C]. "
             "M is the number of bbox, C is the class number. In this case, "
             "Input BBoxes should be the second case with shape [M, C, 4].");
D
dangqingqing 已提交
474
    AddAttr<int>(
475
        "background_label",
翟飞跃 已提交
476
        "(int, default: 0) "
D
dangqingqing 已提交
477 478
        "The index of background label, the background label will be ignored. "
        "If set to -1, then all categories will be considered.")
479
        .SetDefault(0);
D
dangqingqing 已提交
480 481
    AddAttr<float>("score_threshold",
                   "(float) "
D
dangqingqing 已提交
482 483
                   "Threshold to filter out bounding boxes with low "
                   "confidence score. If not provided, consider all boxes.");
D
dangqingqing 已提交
484 485 486
    AddAttr<int>("nms_top_k",
                 "(int64_t) "
                 "Maximum number of detections to be kept according to the "
T
tianshuo78520a 已提交
487
                 "confidences after the filtering detections based on "
D
dangqingqing 已提交
488
                 "score_threshold");
489
    AddAttr<float>("nms_threshold",
翟飞跃 已提交
490
                   "(float, default: 0.3) "
D
dangqingqing 已提交
491
                   "The threshold to be used in NMS.")
492 493 494
        .SetDefault(0.3);
    AddAttr<float>("nms_eta",
                   "(float) "
D
dangqingqing 已提交
495
                   "The parameter for adaptive NMS.")
496
        .SetDefault(1.0);
D
dangqingqing 已提交
497 498 499 500
    AddAttr<int>("keep_top_k",
                 "(int64_t) "
                 "Number of total bboxes to be kept per image after NMS "
                 "step. -1 means keeping all bboxes after NMS step.");
J
jerrywgz 已提交
501
    AddAttr<bool>("normalized",
J
jerrywgz 已提交
502
                  "(bool, default true) "
J
jerrywgz 已提交
503 504
                  "Whether detections are normalized.")
        .SetDefault(true);
505 506 507
    AddOutput("Out",
              "(LoDTensor) A 2-D LoDTensor with shape [No, 6] represents the "
              "detections. Each row has 6 values: "
Y
Yipeng 已提交
508 509 510 511 512 513
              "[label, confidence, xmin, ymin, xmax, ymax] or "
              "(LoDTensor) A 2-D LoDTensor with shape [No, 10] represents the "
              "detections. Each row has 10 values: "
              "[label, confidence, x1, y1, x2, y2, x3, y3, x4, y4]. No is the "
              "total number of detections in this mini-batch."
              "For each instance, "
514 515 516 517
              "the offsets in first dimension are called LoD, the number of "
              "offset is N + 1, if LoD[i + 1] - LoD[i] == 0, means there is "
              "no detected bbox.");
    AddComment(R"DOC(
D
dangqingqing 已提交
518
This operator is to do multi-class non maximum suppression (NMS) on a batched
519
of boxes and scores.
D
dangqingqing 已提交
520 521 522 523 524 525
In the NMS step, this operator greedily selects a subset of detection bounding
boxes that have high scores larger than score_threshold, if providing this
threshold, then selects the largest nms_top_k confidences scores if nms_top_k
is larger than -1. Then this operator pruns away boxes that have high IOU
(intersection over union) overlap with already selected boxes by adaptive
threshold NMS based on parameters of nms_threshold and nms_eta.
526
Aftern NMS step, at most keep_top_k number of total bboxes are to be kept
D
dangqingqing 已提交
527 528
per image if keep_top_k is larger than -1.
This operator support multi-class and batched inputs. It applying NMS
529 530 531
independently for each class. The outputs is a 2-D LoDTenosr, for each
image, the offsets in first dimension of LoDTensor are called LoD, the number
of offset is N + 1, where N is the batch size. If LoD[i + 1] - LoD[i] == 0,
532
means there is no detected bbox for this image.
533 534 535 536
)DOC");
  }
};

537 538 539 540 541 542 543 544 545 546 547 548 549 550
class MultiClassNMS2Op : public MultiClassNMSOp {
 public:
  MultiClassNMS2Op(const std::string& type,
                   const framework::VariableNameMap& inputs,
                   const framework::VariableNameMap& outputs,
                   const framework::AttributeMap& attrs)
      : MultiClassNMSOp(type, inputs, outputs, attrs) {}

  void InferShape(framework::InferShapeContext* ctx) const override {
    MultiClassNMSOp::InferShape(ctx);

    auto score_dims = ctx->GetInputDim("Scores");
    auto score_size = score_dims.size();
    if (score_size == 3) {
551
      ctx->SetOutputDim("Index", {-1, 1});
552 553 554
    } else {
      ctx->SetOutputDim("Index", {-1, 1});
    }
555 556 557
    if (!ctx->IsRuntime()) {
      ctx->SetLoDLevel("Index", std::max(ctx->GetLoDLevel("BBoxes"), 1));
    }
558 559 560 561 562 563 564 565 566 567 568 569 570 571 572
  }
};

class MultiClassNMS2OpMaker : public MultiClassNMSOpMaker {
 public:
  void Make() override {
    MultiClassNMSOpMaker::Make();
    AddOutput("Index",
              "(LoDTensor) A 2-D LoDTensor with shape [No, 1] represents the "
              "index of selected bbox. The index is the absolute index cross "
              "batches.")
        .AsIntermediate();
  }
};

573 574 575 576 577 578 579 580 581 582 583 584 585 586 587 588 589 590 591 592 593 594 595 596 597 598 599 600
class MultiClassNMS3Op : public MultiClassNMS2Op {
 public:
  MultiClassNMS3Op(const std::string& type,
                   const framework::VariableNameMap& inputs,
                   const framework::VariableNameMap& outputs,
                   const framework::AttributeMap& attrs)
      : MultiClassNMS2Op(type, inputs, outputs, attrs) {}

  void InferShape(framework::InferShapeContext* ctx) const override {
    MultiClassNMS2Op::InferShape(ctx);

    ctx->SetOutputDim("NmsRoisNum", {-1});
  }
};

class MultiClassNMS3OpMaker : public MultiClassNMS2OpMaker {
 public:
  void Make() override {
    MultiClassNMS2OpMaker::Make();
    AddInput("RoisNum",
             "(Tensor) The number of RoIs in shape (B),"
             "B is the number of images")
        .AsDispensable();
    AddOutput("NmsRoisNum", "(Tensor), The number of NMS RoIs in each image")
        .AsDispensable();
  }
};

601 602 603 604
}  // namespace operators
}  // namespace paddle

namespace ops = paddle::operators;
H
hong 已提交
605 606 607 608
REGISTER_OPERATOR(
    multiclass_nms, ops::MultiClassNMSOp, ops::MultiClassNMSOpMaker,
    paddle::framework::EmptyGradOpMaker<paddle::framework::OpDesc>,
    paddle::framework::EmptyGradOpMaker<paddle::imperative::OpBase>);
D
dangqingqing 已提交
609 610
REGISTER_OP_CPU_KERNEL(multiclass_nms, ops::MultiClassNMSKernel<float>,
                       ops::MultiClassNMSKernel<double>);
H
hong 已提交
611 612 613 614
REGISTER_OPERATOR(
    multiclass_nms2, ops::MultiClassNMS2Op, ops::MultiClassNMS2OpMaker,
    paddle::framework::EmptyGradOpMaker<paddle::framework::OpDesc>,
    paddle::framework::EmptyGradOpMaker<paddle::imperative::OpBase>);
615 616
REGISTER_OP_CPU_KERNEL(multiclass_nms2, ops::MultiClassNMSKernel<float>,
                       ops::MultiClassNMSKernel<double>);
617 618 619 620 621 622 623

REGISTER_OPERATOR(
    multiclass_nms3, ops::MultiClassNMS3Op, ops::MultiClassNMS3OpMaker,
    paddle::framework::EmptyGradOpMaker<paddle::framework::OpDesc>,
    paddle::framework::EmptyGradOpMaker<paddle::imperative::OpBase>);
REGISTER_OP_CPU_KERNEL(multiclass_nms3, ops::MultiClassNMSKernel<float>,
                       ops::MultiClassNMSKernel<double>);