Merge pull request #22857 from fengyuentau:batched_nms

dnn: add batched nms

modules/dnn/include/opencv2/dnn/dnn.hpp

@@ -1194,6 +1194,27 @@ CV__DNN_INLINE_NS_BEGIN
                              CV_OUT std::vector<int>& indices,
                              const float eta = 1.f, const int top_k = 0);
 
+    /** @brief Performs batched non maximum suppression on given boxes and corresponding scores across different classes.
+
+     * @param bboxes a set of bounding boxes to apply NMS.
+     * @param scores a set of corresponding confidences.
+     * @param class_ids a set of corresponding class ids. Ids are integer and usually start from 0.
+     * @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 eta a coefficient in adaptive threshold formula: \f$nms\_threshold_{i+1}=eta\cdot nms\_threshold_i\f$.
+     * @param top_k if `>0`, keep at most @p top_k picked indices.
+     */
+    CV_EXPORTS void NMSBoxesBatched(const std::vector<Rect>& bboxes, const std::vector<float>& scores, const std::vector<int>& class_ids,
+                                    const float score_threshold, const float nms_threshold,
+                                    CV_OUT std::vector<int>& indices,
+                                    const float eta = 1.f, const int top_k = 0);
+
+    CV_EXPORTS_W void NMSBoxesBatched(const std::vector<Rect2d>& bboxes, const std::vector<float>& scores, const std::vector<int>& class_ids,
+                                      const float score_threshold, const float nms_threshold,
+                                      CV_OUT std::vector<int>& indices,
+                                      const float eta = 1.f, const int top_k = 0);
+
     /**
      * @brief Enum of Soft NMS methods.
      * @see softNMSBoxes

modules/dnn/src/nms.cpp

@@ -58,6 +58,61 @@ void NMSBoxes(const std::vector<RotatedRect>& bboxes, const std::vector<float>&
     NMSFast_(bboxes, scores, score_threshold, nms_threshold, eta, top_k, indices, rotatedRectIOU);
 }
 
+template<class Rect_t>
+static inline void NMSBoxesBatchedImpl(const std::vector<Rect_t>& bboxes,
+                                       const std::vector<float>& scores, const std::vector<int>& class_ids,
+                                       const float score_threshold, const float nms_threshold,
+                                       std::vector<int>& indices, const float eta, const int top_k)
+{
+    double x1, y1, x2, y2, max_coord = 0;
+    for (int i = 0; i < bboxes.size(); i++)
+    {
+        x1 = bboxes[i].x;
+        y1 = bboxes[i].y;
+        x2 = x1 + bboxes[i].width;
+        y2 = y1 + bboxes[i].height;
+
+        max_coord = std::max(x1, max_coord);
+        max_coord = std::max(y1, max_coord);
+        max_coord = std::max(x2, max_coord);
+        max_coord = std::max(y2, max_coord);
+    }
+
+    // calculate offset and add offset to each bbox
+    std::vector<Rect_t> bboxes_offset;
+    double offset;
+    for (int i = 0; i < bboxes.size(); i++)
+    {
+        offset = class_ids[i] * (max_coord + 1);
+        bboxes_offset.push_back(
+            Rect_t(bboxes[i].x + offset, bboxes[i].y + offset,
+                   bboxes[i].width, bboxes[i].height)
+        );
+    }
+
+    NMSFast_(bboxes_offset, scores, score_threshold, nms_threshold, eta, top_k, indices, rectOverlap);
+}
+
+void NMSBoxesBatched(const std::vector<Rect>& bboxes,
+                     const std::vector<float>& scores, const std::vector<int>& class_ids,
+                     const float score_threshold, const float nms_threshold,
+                     std::vector<int>& indices, const float eta, const int top_k)
+{
+    CV_Assert_N(bboxes.size() == scores.size(), scores.size() == class_ids.size(), nms_threshold >= 0, eta > 0);
+
+    NMSBoxesBatchedImpl(bboxes, scores, class_ids, score_threshold, nms_threshold, indices, eta, top_k);
+}
+
+void NMSBoxesBatched(const std::vector<Rect2d>& bboxes,
+                     const std::vector<float>& scores, const std::vector<int>& class_ids,
+                     const float score_threshold, const float nms_threshold,
+                     std::vector<int>& indices, const float eta, const int top_k)
+{
+    CV_Assert_N(bboxes.size() == scores.size(), scores.size() == class_ids.size(), nms_threshold >= 0, eta > 0);
+
+    NMSBoxesBatchedImpl(bboxes, scores, class_ids, score_threshold, nms_threshold, indices, eta, top_k);
+}
+
 void softNMSBoxes(const std::vector<Rect>& bboxes,
                   const std::vector<float>& scores,
                   std::vector<float>& updated_scores,

modules/dnn/test/test_nms.cpp

@@ -37,6 +37,36 @@ TEST(NMS, Accuracy)
         ASSERT_EQ(indices[i], ref_indices[i]);
 }
 
+TEST(BatchedNMS, Accuracy)
+{
+    //reference results obtained using tf.image.non_max_suppression with iou_threshold=0.5
+    std::string dataPath = findDataFile("dnn/batched_nms_reference.yml");
+    FileStorage fs(dataPath, FileStorage::READ);
+
+    std::vector<Rect> bboxes;
+    std::vector<float> scores;
+    std::vector<int> idxs;
+    std::vector<int> ref_indices;
+
+    fs["boxes"] >> bboxes;
+    fs["probs"] >> scores;
+    fs["idxs"] >> idxs;
+    fs["output"] >> ref_indices;
+
+    const float nms_thresh = .5f;
+    const float score_thresh = .05f;
+    std::vector<int> indices;
+    cv::dnn::NMSBoxesBatched(bboxes, scores, idxs, score_thresh, nms_thresh, indices);
+
+    ASSERT_EQ(ref_indices.size(), indices.size());
+
+    std::sort(indices.begin(), indices.end());
+    std::sort(ref_indices.begin(), ref_indices.end());
+
+    for(size_t i = 0; i < indices.size(); i++)
+        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