| 336 | |
| 337 | template <typename T, typename Context> |
| 338 | void MultiClassNMS(const Context& dev_ctx, |
| 339 | const DenseTensor& scores, |
| 340 | const DenseTensor& bboxes, |
| 341 | const int scores_size, |
| 342 | float scorethreshold, |
| 343 | int nms_top_k, |
| 344 | int keep_top_k, |
| 345 | float nmsthreshold, |
| 346 | bool normalized, |
| 347 | float nmseta, |
| 348 | int background_label, |
| 349 | std::map<int, std::vector<int>>* indices, |
| 350 | int* num_nmsed_out) { |
| 351 | T nms_threshold = static_cast<T>(nmsthreshold); |
| 352 | T nms_eta = static_cast<T>(nmseta); |
| 353 | T score_threshold = static_cast<T>(scorethreshold); |
| 354 | |
| 355 | int num_det = 0; |
| 356 | |
| 357 | int class_num = |
| 358 | static_cast<int>(scores_size == 3 ? scores.dims()[0] : scores.dims()[1]); |
| 359 | DenseTensor bbox_slice, score_slice; |
| 360 | for (int c = 0; c < class_num; ++c) { |
| 361 | if (c == background_label) continue; |
| 362 | if (scores_size == 3) { |
| 363 | score_slice = scores.Slice(c, c + 1); |
| 364 | bbox_slice = bboxes; |
| 365 | } else { |
| 366 | score_slice.Resize({scores.dims()[0], 1}); |
| 367 | bbox_slice.Resize({scores.dims()[0], 4}); |
| 368 | SliceOneClass<T, Context>(dev_ctx, scores, c, &score_slice); |
| 369 | SliceOneClass<T, Context>(dev_ctx, bboxes, c, &bbox_slice); |
| 370 | } |
| 371 | NMSFast<T>(bbox_slice, |
| 372 | score_slice, |
| 373 | score_threshold, |
| 374 | nms_threshold, |
| 375 | nms_eta, |
| 376 | nms_top_k, |
| 377 | &((*indices)[c]), |
| 378 | normalized); |
| 379 | if (scores_size == 2) { |
| 380 | std::stable_sort((*indices)[c].begin(), (*indices)[c].end()); |
| 381 | } |
| 382 | num_det += static_cast<int>((*indices)[c].size()); |
| 383 | } |
| 384 | |
| 385 | *num_nmsed_out = num_det; |
| 386 | const T* scores_data = scores.data<T>(); |
| 387 | if (keep_top_k > -1 && num_det > keep_top_k) { |
| 388 | const T* sdata = nullptr; |
| 389 | std::vector<std::pair<float, std::pair<int, int>>> score_index_pairs; |
| 390 | for (const auto& it : *indices) { |
| 391 | int label = it.first; |
| 392 | if (scores_size == 3) { |
| 393 | sdata = scores_data + label * scores.dims()[1]; |
| 394 | } else { |
| 395 | score_slice.Resize({scores.dims()[0], 1}); |