diff --git a/modules/dnn/include/opencv2/dnn/dnn.hpp b/modules/dnn/include/opencv2/dnn/dnn.hpp index bf1670051a..8b4e5e21d4 100644 --- a/modules/dnn/include/opencv2/dnn/dnn.hpp +++ b/modules/dnn/include/opencv2/dnn/dnn.hpp @@ -1130,6 +1130,39 @@ CV__DNN_INLINE_NS_BEGIN CV_OUT std::vector& indices, const float eta = 1.f, const int top_k = 0); + /** + * @brief Enum of Soft NMS methods. + * @see softNMSBoxes + */ + enum class SoftNMSMethod + { + SOFTNMS_LINEAR = 1, + SOFTNMS_GAUSSIAN = 2 + }; + + /** @brief Performs soft non maximum suppression given boxes and corresponding scores. + * Reference: https://arxiv.org/abs/1704.04503 + * @param bboxes a set of bounding boxes to apply Soft NMS. + * @param scores a set of corresponding confidences. + * @param updated_scores a set of corresponding updated confidences. + * @param score_threshold a threshold used to filter boxes by score. + * @param nms_threshold a threshold used in non maximum suppression. + * @param indices the kept indices of bboxes after NMS. + * @param top_k keep at most @p top_k picked indices. + * @param sigma parameter of Gaussian weighting. + * @param method Gaussian or linear. + * @see SoftNMSMethod + */ + CV_EXPORTS_W void softNMSBoxes(const std::vector& bboxes, + const std::vector& scores, + CV_OUT std::vector& updated_scores, + const float score_threshold, + const float nms_threshold, + CV_OUT std::vector& indices, + size_t top_k = 0, + const float sigma = 0.5, + SoftNMSMethod method = SoftNMSMethod::SOFTNMS_GAUSSIAN); + /** @brief This class is presented high-level API for neural networks. * diff --git a/modules/dnn/src/nms.cpp b/modules/dnn/src/nms.cpp index 26489dc7e9..d321cfed88 100644 --- a/modules/dnn/src/nms.cpp +++ b/modules/dnn/src/nms.cpp @@ -58,6 +58,83 @@ void NMSBoxes(const std::vector& bboxes, const std::vector& NMSFast_(bboxes, scores, score_threshold, nms_threshold, eta, top_k, indices, rotatedRectIOU); } +void softNMSBoxes(const std::vector& bboxes, + const std::vector& scores, + std::vector& updated_scores, + const float score_threshold, + const float nms_threshold, + std::vector& indices, + size_t top_k, + const float sigma, + SoftNMSMethod method) +{ + CV_Assert_N(bboxes.size() == scores.size(), score_threshold >= 0, + nms_threshold >= 0, sigma >= 0); + + indices.clear(); + updated_scores.clear(); + + std::vector > score_index_vec(scores.size()); + for (size_t i = 0; i < scores.size(); i++) + { + score_index_vec[i].first = scores[i]; + score_index_vec[i].second = i; + } + + const auto score_cmp = [](const std::pair& a, const std::pair& b) + { + return a.first == b.first ? a.second > b.second : a.first < b.first; + }; + + top_k = top_k == 0 ? scores.size() : std::min(top_k, scores.size()); + ptrdiff_t start = 0; + while (indices.size() < top_k) + { + auto it = std::max_element(score_index_vec.begin() + start, score_index_vec.end(), score_cmp); + + float bscore = it->first; + size_t bidx = it->second; + + if (bscore < score_threshold) + { + break; + } + + indices.push_back(static_cast(bidx)); + updated_scores.push_back(bscore); + std::swap(score_index_vec[start], *it); // first start elements are chosen + + for (size_t i = start + 1; i < scores.size(); ++i) + { + float& bscore_i = score_index_vec[i].first; + const size_t bidx_i = score_index_vec[i].second; + + if (bscore_i < score_threshold) + { + continue; + } + + float overlap = rectOverlap(bboxes[bidx], bboxes[bidx_i]); + + switch (method) + { + case SoftNMSMethod::SOFTNMS_LINEAR: + if (overlap > nms_threshold) + { + bscore_i *= 1.f - overlap; + } + break; + case SoftNMSMethod::SOFTNMS_GAUSSIAN: + bscore_i *= exp(-(overlap * overlap) / sigma); + break; + default: + CV_Error(Error::StsBadArg, "Not supported SoftNMS method."); + } + } + ++start; + } +} + CV__DNN_INLINE_NS_END }// dnn }// cv diff --git a/modules/dnn/test/test_nms.cpp b/modules/dnn/test/test_nms.cpp index 53cb3500ec..5ba240287c 100644 --- a/modules/dnn/test/test_nms.cpp +++ b/modules/dnn/test/test_nms.cpp @@ -37,4 +37,41 @@ TEST(NMS, Accuracy) ASSERT_EQ(indices[i], ref_indices[i]); } +TEST(SoftNMS, Accuracy) +{ + //reference results are obtained using TF v2.7 tf.image.non_max_suppression_with_scores + std::string dataPath = findDataFile("dnn/soft_nms_reference.yml"); + FileStorage fs(dataPath, FileStorage::READ); + + std::vector bboxes; + std::vector scores; + std::vector ref_indices; + std::vector ref_updated_scores; + + fs["boxes"] >> bboxes; + fs["probs"] >> scores; + fs["indices"] >> ref_indices; + fs["updated_scores"] >> ref_updated_scores; + + std::vector updated_scores; + const float score_thresh = .01f; + const float nms_thresh = .5f; + std::vector indices; + const size_t top_k = 0; + const float sigma = 1.; // sigma in TF is being multiplied by 2, so 0.5 should be passed there + cv::dnn::softNMSBoxes(bboxes, scores, updated_scores, score_thresh, nms_thresh, indices, top_k, sigma); + + ASSERT_EQ(ref_indices.size(), indices.size()); + for(size_t i = 0; i < indices.size(); i++) + { + ASSERT_EQ(indices[i], ref_indices[i]); + } + + ASSERT_EQ(ref_updated_scores.size(), updated_scores.size()); + for(size_t i = 0; i < updated_scores.size(); i++) + { + EXPECT_NEAR(updated_scores[i], ref_updated_scores[i], 1e-7); + } +} + }} // namespace