multiclass_nms_op.cc 23.7 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"
Y
Yipeng 已提交
16
#include "paddle/fluid/operators/detection/poly_util.h"
17 18 19 20 21 22 23

namespace paddle {
namespace operators {

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

D
dangqingqing 已提交
24
class MultiClassNMSOp : public framework::OperatorWithKernel {
25 26 27 28
 public:
  using framework::OperatorWithKernel::OperatorWithKernel;

  void InferShape(framework::InferShapeContext* ctx) const override {
D
dangqingqing 已提交
29 30
    PADDLE_ENFORCE(ctx->HasInput("BBoxes"),
                   "Input(BBoxes) of MultiClassNMS should not be null.");
31
    PADDLE_ENFORCE(ctx->HasInput("Scores"),
D
dangqingqing 已提交
32 33 34
                   "Input(Scores) of MultiClassNMS should not be null.");
    PADDLE_ENFORCE(ctx->HasOutput("Out"),
                   "Output(Out) of MultiClassNMS should not be null.");
35

D
dangqingqing 已提交
36
    auto box_dims = ctx->GetInputDim("BBoxes");
37
    auto score_dims = ctx->GetInputDim("Scores");
J
jerrywgz 已提交
38
    auto score_size = score_dims.size();
39

40
    if (ctx->IsRuntime()) {
J
jerrywgz 已提交
41 42
      PADDLE_ENFORCE(score_size == 2 || score_size == 3,
                     "The rank of Input(Scores) must be 2 or 3");
43
      PADDLE_ENFORCE_EQ(box_dims.size(), 3,
J
jerrywgz 已提交
44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68
                        "The rank of Input(BBoxes) must be 3");
      if (score_size == 3) {
        PADDLE_ENFORCE(box_dims[2] == 4 || box_dims[2] == 8 ||
                           box_dims[2] == 16 || box_dims[2] == 24 ||
                           box_dims[2] == 32,
                       "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");
        PADDLE_ENFORCE_EQ(
            box_dims[1], score_dims[2],
            "The 2nd dimension of Input(BBoxes) must be equal to "
            "last dimension of Input(Scores), which represents the "
            "predicted bboxes.");
      } else {
        PADDLE_ENFORCE(box_dims[2] == 4,
                       "The last dimension of Input(BBoxes) must be 4");
        PADDLE_ENFORCE_EQ(box_dims[1], score_dims[1],
                          "The 2nd dimension of Input(BBoxes)"
                          "must be equal to the 2nd dimension"
                          " of Input(Scores)");
      }
69
    }
70 71
    // Here the box_dims[0] is not the real dimension of output.
    // It will be rewritten in the computing kernel.
J
jerrywgz 已提交
72 73 74 75 76
    if (score_size == 3) {
      ctx->SetOutputDim("Out", {box_dims[1], box_dims[2] + 2});
    } else {
      ctx->SetOutputDim("Out", {-1, box_dims[2] + 2});
    }
77
  }
D
dangqingqing 已提交
78 79 80 81 82

 protected:
  framework::OpKernelType GetExpectedKernelType(
      const framework::ExecutionContext& ctx) const override {
    return framework::OpKernelType(
83
        OperatorWithKernel::IndicateVarDataType(ctx, "Scores"),
84
        platform::CPUPlace());
D
dangqingqing 已提交
85
  }
86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106
};

template <class T>
bool SortScorePairDescend(const std::pair<float, T>& pair1,
                          const std::pair<float, T>& pair2) {
  return pair1.first > pair2.first;
}

template <class T>
static inline void GetMaxScoreIndex(
    const std::vector<T>& scores, const T threshold, int top_k,
    std::vector<std::pair<T, int>>* sorted_indices) {
  for (size_t i = 0; i < scores.size(); ++i) {
    if (scores[i] > threshold) {
      sorted_indices->push_back(std::make_pair(scores[i], i));
    }
  }
  // Sort the score pair according to the scores in descending order
  std::stable_sort(sorted_indices->begin(), sorted_indices->end(),
                   SortScorePairDescend<int>);
  // Keep top_k scores if needed.
107
  if (top_k > -1 && top_k < static_cast<int>(sorted_indices->size())) {
108 109 110 111 112
    sorted_indices->resize(top_k);
  }
}

template <class T>
113
static inline T BBoxArea(const T* box, const bool normalized) {
114
  if (box[2] < box[0] || box[3] < box[1]) {
D
dangqingqing 已提交
115 116 117
    // If coordinate values are is invalid
    // (e.g. xmax < xmin or ymax < ymin), return 0.
    return static_cast<T>(0.);
118 119 120 121 122 123
  } else {
    const T w = box[2] - box[0];
    const T h = box[3] - box[1];
    if (normalized) {
      return w * h;
    } else {
D
dangqingqing 已提交
124
      // If coordinate values are not within range [0, 1].
125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140
      return (w + 1) * (h + 1);
    }
  }
}

template <class T>
static inline T JaccardOverlap(const T* box1, const T* box2,
                               const bool normalized) {
  if (box2[0] > box1[2] || box2[2] < box1[0] || box2[1] > box1[3] ||
      box2[3] < box1[1]) {
    return static_cast<T>(0.);
  } else {
    const T inter_xmin = std::max(box1[0], box2[0]);
    const T inter_ymin = std::max(box1[1], box2[1]);
    const T inter_xmax = std::min(box1[2], box2[2]);
    const T inter_ymax = std::min(box1[3], box2[3]);
141 142 143
    T norm = normalized ? static_cast<T>(0.) : static_cast<T>(1.);
    T inter_w = inter_xmax - inter_xmin + norm;
    T inter_h = inter_ymax - inter_ymin + norm;
144 145 146 147 148 149 150
    const T inter_area = inter_w * inter_h;
    const T bbox1_area = BBoxArea<T>(box1, normalized);
    const T bbox2_area = BBoxArea<T>(box2, normalized);
    return inter_area / (bbox1_area + bbox2_area - inter_area);
  }
}

Y
Yipeng 已提交
151 152 153 154 155 156 157
template <class T>
T PolyIoU(const T* box1, const T* box2, const size_t box_size,
          const bool normalized) {
  T bbox1_area = PolyArea<T>(box1, box_size, normalized);
  T bbox2_area = PolyArea<T>(box2, box_size, normalized);
  T inter_area = PolyOverlapArea<T>(box1, box2, box_size, normalized);
  if (bbox1_area == 0 || bbox2_area == 0 || inter_area == 0) {
J
jerrywgz 已提交
158
    // If coordinate values are invalid
Y
Yipeng 已提交
159 160 161 162 163 164 165
    // if area size <= 0,  return 0.
    return T(0.);
  } else {
    return inter_area / (bbox1_area + bbox2_area - inter_area);
  }
}

166 167 168 169 170 171 172 173 174
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 已提交
175 176 177 178 179 180 181 182 183 184
    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];
    }
185 186 187
  }
}

188
template <typename T>
D
dangqingqing 已提交
189
class MultiClassNMSKernel : public framework::OpKernel<T> {
190 191 192
 public:
  void NMSFast(const Tensor& bbox, const Tensor& scores,
               const T score_threshold, const T nms_threshold, const T eta,
J
jerrywgz 已提交
193 194
               const int64_t top_k, std::vector<int>* selected_indices,
               const bool normalized) const {
195 196 197
    // The total boxes for each instance.
    int64_t num_boxes = bbox.dims()[0];
    // 4: [xmin ymin xmax ymax]
Y
Yipeng 已提交
198 199
    // 8: [x1 y1 x2 y2 x3 y3 x4 y4]
    // 16, 24, or 32: [x1 y1 x2 y2 ...  xn yn], n = 8, 12 or 16
200 201 202 203 204 205 206 207 208 209 210 211 212 213
    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;
214
      for (size_t k = 0; k < selected_indices->size(); ++k) {
215 216
        if (keep) {
          const int kept_idx = (*selected_indices)[k];
Y
Yipeng 已提交
217 218 219
          T overlap = T(0.);
          // 4: [xmin ymin xmax ymax]
          if (box_size == 4) {
J
jerrywgz 已提交
220 221 222
            overlap =
                JaccardOverlap<T>(bbox_data + idx * box_size,
                                  bbox_data + kept_idx * box_size, normalized);
Y
Yipeng 已提交
223 224 225 226
          }
          // 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 已提交
227 228 229
            overlap = PolyIoU<T>(bbox_data + idx * box_size,
                                 bbox_data + kept_idx * box_size, box_size,
                                 normalized);
Y
Yipeng 已提交
230
          }
231 232 233 234 235 236 237 238 239 240 241 242 243 244 245
          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 已提交
246
  void MultiClassNMS(const framework::ExecutionContext& ctx,
247
                     const Tensor& scores, const Tensor& bboxes,
J
jerrywgz 已提交
248
                     const int scores_size,
249 250
                     std::map<int, std::vector<int>>* indices,
                     int* num_nmsed_out) const {
D
dangqingqing 已提交
251 252 253
    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 已提交
254
    bool normalized = ctx.Attr<bool>("normalized");
255 256
    T nms_threshold = static_cast<T>(ctx.Attr<float>("nms_threshold"));
    T nms_eta = static_cast<T>(ctx.Attr<float>("nms_eta"));
D
dangqingqing 已提交
257
    T score_threshold = static_cast<T>(ctx.Attr<float>("score_threshold"));
J
jerrywgz 已提交
258
    auto& dev_ctx = ctx.template device_context<platform::CPUDeviceContext>();
259 260

    int num_det = 0;
261 262 263 264 265 266 267 268 269 270 271 272 273

    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 已提交
274
      }
275 276 277
      NMSFast(bbox_slice, score_slice, score_threshold, nms_threshold, nms_eta,
              nms_top_k, &((*indices)[c]), normalized);
      if (scores_size == 2) {
J
jerrywgz 已提交
278 279
        std::stable_sort((*indices)[c].begin(), (*indices)[c].end());
      }
280
      num_det += (*indices)[c].size();
281 282
    }

283
    *num_nmsed_out = num_det;
284 285
    const T* scores_data = scores.data<T>();
    if (keep_top_k > -1 && num_det > keep_top_k) {
J
jerrywgz 已提交
286
      const T* sdata;
287
      std::vector<std::pair<float, std::pair<int, int>>> score_index_pairs;
288
      for (const auto& it : *indices) {
289
        int label = it.first;
J
jerrywgz 已提交
290
        if (scores_size == 3) {
291
          sdata = scores_data + label * scores.dims()[1];
J
jerrywgz 已提交
292
        } else {
293 294 295
          score_slice.Resize({scores.dims()[0], 1});
          SliceOneClass<T>(dev_ctx, scores, label, &score_slice);
          sdata = score_slice.data<T>();
J
jerrywgz 已提交
296
        }
297
        const std::vector<int>& label_indices = it.second;
298
        for (size_t j = 0; j < label_indices.size(); ++j) {
299 300 301 302 303 304
          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.
305 306
      std::stable_sort(score_index_pairs.begin(), score_index_pairs.end(),
                       SortScorePairDescend<std::pair<int, int>>);
307 308 309 310
      score_index_pairs.resize(keep_top_k);

      // Store the new indices.
      std::map<int, std::vector<int>> new_indices;
311
      for (size_t j = 0; j < score_index_pairs.size(); ++j) {
312 313 314 315
        int label = score_index_pairs[j].second.first;
        int idx = score_index_pairs[j].second.second;
        new_indices[label].push_back(idx);
      }
J
jerrywgz 已提交
316 317 318 319 320 321 322
      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());
        }
      }
323 324
      new_indices.swap(*indices);
      *num_nmsed_out = keep_top_k;
325 326 327
    }
  }

J
jerrywgz 已提交
328 329
  void MultiClassOutput(const platform::DeviceContext& ctx,
                        const Tensor& scores, const Tensor& bboxes,
330
                        const std::map<int, std::vector<int>>& selected_indices,
331 332
                        const int scores_size, Tensor* outs,
                        int* oindices = nullptr, const int offset = 0) const {
J
jerrywgz 已提交
333
    int64_t class_num = scores.dims()[1];
Y
Yipeng 已提交
334 335
    int64_t predict_dim = scores.dims()[1];
    int64_t box_size = bboxes.dims()[1];
J
jerrywgz 已提交
336 337 338 339
    if (scores_size == 2) {
      box_size = bboxes.dims()[2];
    }
    int64_t out_dim = box_size + 2;
340 341 342
    auto* scores_data = scores.data<T>();
    auto* bboxes_data = bboxes.data<T>();
    auto* odata = outs->data<T>();
J
jerrywgz 已提交
343 344 345
    const T* sdata;
    Tensor bbox;
    bbox.Resize({scores.dims()[0], box_size});
346 347 348
    int count = 0;
    for (const auto& it : selected_indices) {
      int label = it.first;
D
dangqingqing 已提交
349
      const std::vector<int>& indices = it.second;
J
jerrywgz 已提交
350 351 352 353 354
      if (scores_size == 2) {
        SliceOneClass<T>(ctx, bboxes, label, &bbox);
      } else {
        sdata = scores_data + label * predict_dim;
      }
355
      for (size_t j = 0; j < indices.size(); ++j) {
356
        int idx = indices[j];
J
jerrywgz 已提交
357 358 359 360 361
        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
362 363 364
          if (oindices != nullptr) {
            oindices[count] = offset + idx;
          }
J
jerrywgz 已提交
365 366 367
        } else {
          bdata = bbox.data<T>() + idx * box_size;
          odata[count * out_dim + 1] = *(scores_data + idx * class_num + label);
368 369 370
          if (oindices != nullptr) {
            oindices[count] = offset + idx * class_num + label;
          }
J
jerrywgz 已提交
371
        }
Y
Yipeng 已提交
372 373
        // xmin, ymin, xmax, ymax or multi-points coordinates
        std::memcpy(odata + count * out_dim + 2, bdata, box_size * sizeof(T));
D
dangqingqing 已提交
374
        count++;
375 376 377 378 379
      }
    }
  }

  void Compute(const framework::ExecutionContext& ctx) const override {
J
jerrywgz 已提交
380 381
    auto* boxes = ctx.Input<LoDTensor>("BBoxes");
    auto* scores = ctx.Input<LoDTensor>("Scores");
382
    auto* outs = ctx.Output<LoDTensor>("Out");
383 384
    bool return_index = ctx.HasOutput("Index") ? true : false;
    auto index = ctx.Output<LoDTensor>("Index");
385
    auto score_dims = scores->dims();
386
    auto score_size = score_dims.size();
J
jerrywgz 已提交
387
    auto& dev_ctx = ctx.template device_context<platform::CPUDeviceContext>();
388 389 390

    std::vector<std::map<int, std::vector<int>>> all_indices;
    std::vector<size_t> batch_starts = {0};
J
jerrywgz 已提交
391 392 393 394
    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;
395 396 397 398 399 400 401 402 403 404 405 406
    Tensor boxes_slice, scores_slice;
    int n = score_size == 3 ? batch_size : boxes->lod().back().size() - 1;
    for (int i = 0; i < n; ++i) {
      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 {
        auto boxes_lod = boxes->lod().back();
        scores_slice = scores->Slice(boxes_lod[i], boxes_lod[i + 1]);
        boxes_slice = boxes->Slice(boxes_lod[i], boxes_lod[i + 1]);
J
jerrywgz 已提交
407
      }
408 409 410 411 412
      std::map<int, std::vector<int>> indices;
      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 已提交
413 414 415 416
    }

    int num_kept = batch_starts.back();
    if (num_kept == 0) {
417 418 419 420 421 422 423 424
      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 已提交
425 426
    } else {
      outs->mutable_data<T>({num_kept, out_dim}, ctx.GetPlace());
427 428
      int offset = 0;
      int* oindices = nullptr;
429 430 431 432 433 434
      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});
435 436 437
          if (return_index) {
            offset = i * score_dims[2];
          }
438 439 440 441
        } else {
          auto boxes_lod = boxes->lod().back();
          scores_slice = scores->Slice(boxes_lod[i], boxes_lod[i + 1]);
          boxes_slice = boxes->Slice(boxes_lod[i], boxes_lod[i + 1]);
442 443 444
          if (return_index) {
            offset = boxes_lod[i] * score_dims[1];
          }
J
jerrywgz 已提交
445
        }
446 447 448 449
        int64_t s = batch_starts[i];
        int64_t e = batch_starts[i + 1];
        if (e > s) {
          Tensor out = outs->Slice(s, e);
450 451 452 453 454
          if (return_index) {
            int* output_idx =
                index->mutable_data<int>({num_kept, 1}, ctx.GetPlace());
            oindices = output_idx + s;
          }
455
          MultiClassOutput(dev_ctx, scores_slice, boxes_slice, all_indices[i],
456
                           score_dims.size(), &out, oindices, offset);
457 458 459 460 461 462
        }
      }
    }

    framework::LoD lod;
    lod.emplace_back(batch_starts);
463 464 465
    if (return_index) {
      index->set_lod(lod);
    }
466 467 468 469
    outs->set_lod(lod);
  }
};

D
dangqingqing 已提交
470
class MultiClassNMSOpMaker : public framework::OpProtoAndCheckerMaker {
471
 public:
Y
Yu Yang 已提交
472
  void Make() override {
D
dangqingqing 已提交
473
    AddInput("BBoxes",
J
jerrywgz 已提交
474 475
             "Two types of bboxes are supported:"
             "1. (Tensor) A 3-D Tensor with shape "
Y
Yipeng 已提交
476
             "[N, M, 4 or 8 16 24 32] represents the "
477 478
             "predicted locations of M bounding bboxes, N is the batch size. "
             "Each bounding box has four coordinate values and the layout is "
J
jerrywgz 已提交
479
             "[xmin, ymin, xmax, ymax], when box size equals to 4."
480 481
             "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 已提交
482
    AddInput("Scores",
J
jerrywgz 已提交
483 484
             "Two types of scores are supported:"
             "1. (Tensor) A 3-D Tensor with shape [N, C, M] represents the "
D
dangqingqing 已提交
485 486 487
             "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. "
488 489 490 491
             " 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 已提交
492
    AddAttr<int>(
493
        "background_label",
翟飞跃 已提交
494
        "(int, default: 0) "
D
dangqingqing 已提交
495 496
        "The index of background label, the background label will be ignored. "
        "If set to -1, then all categories will be considered.")
497
        .SetDefault(0);
D
dangqingqing 已提交
498 499
    AddAttr<float>("score_threshold",
                   "(float) "
D
dangqingqing 已提交
500 501
                   "Threshold to filter out bounding boxes with low "
                   "confidence score. If not provided, consider all boxes.");
D
dangqingqing 已提交
502 503 504 505 506
    AddAttr<int>("nms_top_k",
                 "(int64_t) "
                 "Maximum number of detections to be kept according to the "
                 "confidences aftern the filtering detections based on "
                 "score_threshold");
507
    AddAttr<float>("nms_threshold",
翟飞跃 已提交
508
                   "(float, default: 0.3) "
D
dangqingqing 已提交
509
                   "The threshold to be used in NMS.")
510 511 512
        .SetDefault(0.3);
    AddAttr<float>("nms_eta",
                   "(float) "
D
dangqingqing 已提交
513
                   "The parameter for adaptive NMS.")
514
        .SetDefault(1.0);
D
dangqingqing 已提交
515 516 517 518
    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 已提交
519
    AddAttr<bool>("normalized",
J
jerrywgz 已提交
520
                  "(bool, default true) "
J
jerrywgz 已提交
521 522
                  "Whether detections are normalized.")
        .SetDefault(true);
523 524 525
    AddOutput("Out",
              "(LoDTensor) A 2-D LoDTensor with shape [No, 6] represents the "
              "detections. Each row has 6 values: "
Y
Yipeng 已提交
526 527 528 529 530 531
              "[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, "
532 533 534 535
              "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 已提交
536
This operator is to do multi-class non maximum suppression (NMS) on a batched
537
of boxes and scores.
D
dangqingqing 已提交
538 539 540 541 542 543
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.
544
Aftern NMS step, at most keep_top_k number of total bboxes are to be kept
D
dangqingqing 已提交
545 546
per image if keep_top_k is larger than -1.
This operator support multi-class and batched inputs. It applying NMS
547 548 549
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,
550
means there is no detected bbox for this image.
551 552 553 554
)DOC");
  }
};

555 556 557 558 559 560 561 562 563 564 565 566 567 568 569 570 571 572 573 574 575 576 577 578 579 580 581 582 583 584 585 586 587 588
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 box_dims = ctx->GetInputDim("BBoxes");
    auto score_dims = ctx->GetInputDim("Scores");
    auto score_size = score_dims.size();
    if (score_size == 3) {
      ctx->SetOutputDim("Index", {box_dims[1], 1});
    } else {
      ctx->SetOutputDim("Index", {-1, 1});
    }
  }
};

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();
  }
};

589 590 591 592
}  // namespace operators
}  // namespace paddle

namespace ops = paddle::operators;
H
hong 已提交
593 594 595 596
REGISTER_OPERATOR(
    multiclass_nms, ops::MultiClassNMSOp, ops::MultiClassNMSOpMaker,
    paddle::framework::EmptyGradOpMaker<paddle::framework::OpDesc>,
    paddle::framework::EmptyGradOpMaker<paddle::imperative::OpBase>);
D
dangqingqing 已提交
597 598
REGISTER_OP_CPU_KERNEL(multiclass_nms, ops::MultiClassNMSKernel<float>,
                       ops::MultiClassNMSKernel<double>);
H
hong 已提交
599 600 601 602
REGISTER_OPERATOR(
    multiclass_nms2, ops::MultiClassNMS2Op, ops::MultiClassNMS2OpMaker,
    paddle::framework::EmptyGradOpMaker<paddle::framework::OpDesc>,
    paddle::framework::EmptyGradOpMaker<paddle::imperative::OpBase>);
603 604
REGISTER_OP_CPU_KERNEL(multiclass_nms2, ops::MultiClassNMSKernel<float>,
                       ops::MultiClassNMSKernel<double>);