Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

add dnn nms #119

Merged
merged 10 commits into from
Jun 1, 2024
Merged
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Next Next commit
add dnn nms
nihui committed May 30, 2024
commit 0973769573ccba9ec72b0abca8b0ff6b40ce8720
7 changes: 6 additions & 1 deletion highgui/CMakeLists.txt
Original file line number Diff line number Diff line change
@@ -10,6 +10,9 @@ set(highgui_srcs
${CMAKE_CURRENT_LIST_DIR}/src/highgui.cpp
${CMAKE_CURRENT_LIST_DIR}/src/kanna_rotate.cpp
${CMAKE_CURRENT_LIST_DIR}/src/videocapture.cpp

# dnn
${CMAKE_CURRENT_LIST_DIR}/src/nms.cpp
)

if(WITH_CVI)
@@ -40,7 +43,9 @@ endif()
file(GLOB highgui_ext_hdrs
"${CMAKE_CURRENT_LIST_DIR}/include/opencv2/*.hpp"
"${CMAKE_CURRENT_LIST_DIR}/include/opencv2/${name}/*.hpp"
"${CMAKE_CURRENT_LIST_DIR}/include/opencv2/${name}/*.h")
"${CMAKE_CURRENT_LIST_DIR}/include/opencv2/${name}/*.h"
"${CMAKE_CURRENT_LIST_DIR}/include/opencv2/dnn/*.hpp"
"${CMAKE_CURRENT_LIST_DIR}/include/opencv2/dnn/*.h")

if(UNIX OR OPENCV_VERSION_MAJOR GREATER_EQUAL 3)
#these variables are set by CHECK_MODULE macro
17 changes: 17 additions & 0 deletions highgui/include/opencv2/dnn.hpp
Original file line number Diff line number Diff line change
@@ -0,0 +1,17 @@
//
// Copyright (C) 2024 nihui
//
// 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
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// 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.
// See the License for the specific language governing permissions and
// limitations under the License.
//

#include "opencv2/dnn/dnn.hpp"
55 changes: 55 additions & 0 deletions highgui/include/opencv2/dnn/dnn.hpp
Original file line number Diff line number Diff line change
@@ -0,0 +1,55 @@
//
// Copyright (C) 2024 nihui
//
// 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
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// 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.
// See the License for the specific language governing permissions and
// limitations under the License.
//

#ifndef OPENCV_DNN_HPP
#define OPENCV_DNN_HPP

#include "opencv2/core.hpp"

namespace cv {
namespace dnn {

enum SoftNMSMethod
{
SOFTNMS_LINEAR = 1,
SOFTNMS_GAUSSIAN = 2
};

CV_EXPORTS void NMSBoxes(const std::vector<Rect>& bboxes, const std::vector<float>& scores,
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_AS(NMSBoxesRotated) void NMSBoxes(const std::vector<RotatedRect>& bboxes, const std::vector<float>& scores,
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 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 softNMSBoxes(const std::vector<Rect>& bboxes, const std::vector<float>& scores,
CV_OUT std::vector<float>& updated_scores,
const float score_threshold, const float nms_threshold,
CV_OUT std::vector<int>& indices,
size_t top_k = 0, const float sigma = 0.5, SoftNMSMethod method = SOFTNMS_GAUSSIAN);

} // namespace dnn
} // namespace cv

#endif // OPENCV_DNN_HPP
259 changes: 259 additions & 0 deletions highgui/src/nms.cpp
Original file line number Diff line number Diff line change
@@ -0,0 +1,259 @@
//
// Copyright (C) 2024 nihui
//
// 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
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// 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.
// See the License for the specific language governing permissions and
// limitations under the License.
//

#include <opencv2/core.hpp>
#include <opencv2/imgproc.hpp>
#include <opencv2/dnn.hpp>
#include <limits>
#include <vector>
#include <algorithm>

namespace cv {
namespace dnn {

static inline bool SortScorePairDescend(const std::pair<float, int>& pair1, const std::pair<float, int>& pair2)
{
return pair1.first > pair2.first;
}

// Get max scores with corresponding indices.
// scores: a set of scores.
// threshold: only consider scores higher than the threshold.
// top_k: if -1, keep all; otherwise, keep at most top_k.
// score_index_vec: store the sorted (score, index) pair.
inline void GetMaxScoreIndex(const std::vector<float>& scores, const float threshold, const int top_k,
std::vector<std::pair<float, int> >& score_index_vec)
{
CV_DbgAssert(score_index_vec.empty());
// Generate index score pairs.
for (size_t i = 0; i < scores.size(); ++i)
{
if (scores[i] > threshold)
{
score_index_vec.push_back(std::make_pair(scores[i], i));
}
}

// Sort the score pair according to the scores in descending order
std::stable_sort(score_index_vec.begin(), score_index_vec.end(), SortScorePairDescend);

// Keep top_k scores if needed.
if (top_k > 0 && top_k < (int)score_index_vec.size())
{
score_index_vec.resize(top_k);
}
}

// Do non maximum suppression given bboxes and scores.
// Inspired by Piotr Dollar's NMS implementation in EdgeBox.
// https://goo.gl/jV3JYS
// bboxes: a set of bounding boxes.
// scores: a set of corresponding confidences.
// score_threshold: a threshold used to filter detection results.
// nms_threshold: a threshold used in non maximum suppression.
// top_k: if not > 0, keep at most top_k picked indices.
// limit: early terminate once the # of picked indices has reached it.
// indices: the kept indices of bboxes after nms.
template <typename BoxType>
inline void NMSFast_(const std::vector<BoxType>& bboxes,
const std::vector<float>& scores, const float score_threshold,
const float nms_threshold, const float eta, const int top_k,
std::vector<int>& indices,
float (*computeOverlap)(const BoxType&, const BoxType&),
int limit = std::numeric_limits<int>::max())
{
CV_Assert(bboxes.size() == scores.size());

// Get top_k scores (with corresponding indices).
std::vector<std::pair<float, int> > score_index_vec;
GetMaxScoreIndex(scores, score_threshold, top_k, score_index_vec);

// Do nms.
float adaptive_threshold = nms_threshold;
indices.clear();
for (size_t i = 0; i < score_index_vec.size(); ++i) {
const int idx = score_index_vec[i].second;
bool keep = true;
for (int k = 0; k < (int)indices.size() && keep; ++k) {
const int kept_idx = indices[k];
float overlap = computeOverlap(bboxes[idx], bboxes[kept_idx]);
keep = overlap <= adaptive_threshold;
}
if (keep) {
indices.push_back(idx);
if ((int)indices.size() >= limit) {
break;
}
}
if (keep && eta < 1 && adaptive_threshold > 0.5) {
adaptive_threshold *= eta;
}
}
}

static inline float rectOverlap(const Rect& a, const Rect& b)
{
return 1.f - static_cast<float>(jaccardDistance(a, b));
}

void NMSBoxes(const std::vector<Rect>& bboxes, const std::vector<float>& scores,
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(), score_threshold >= 0, nms_threshold >= 0, eta > 0);

NMSFast_(bboxes, scores, score_threshold, nms_threshold, eta, top_k, indices, rectOverlap);
}

static inline float rotatedRectIOU(const RotatedRect& a, const RotatedRect& b)
{
std::vector<Point2f> inter;
int res = rotatedRectangleIntersection(a, b, inter);
if (inter.empty() || res == INTERSECT_NONE)
return 0.0f;
if (res == INTERSECT_FULL)
return 1.0f;
float interArea = contourArea(inter);
return interArea / (a.size.area() + b.size.area() - interArea);
}

void NMSBoxes(const std::vector<RotatedRect>& bboxes, const std::vector<float>& scores,
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(), score_threshold >= 0, nms_threshold >= 0, eta > 0);

NMSFast_(bboxes, scores, score_threshold, nms_threshold, eta, top_k, indices, rotatedRectIOU);
}

static inline void NMSBoxesBatchedImpl(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)
{
int x1, y1, x2, y2, max_coord = 0;
for (size_t 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> bboxes_offset;
for (size_t i = 0; i < bboxes.size(); i++)
{
int offset = class_ids[i] * (max_coord + 1);
bboxes_offset.push_back(Rect(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 softNMSBoxes(const std::vector<Rect>& bboxes,
const std::vector<float>& scores,
std::vector<float>& updated_scores,
const float score_threshold,
const float nms_threshold,
std::vector<int>& 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<std::pair<float, size_t> > 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<float, size_t>& a, const std::pair<float, size_t>& 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<int>(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;
}
}

} // namespace dnn
} // namespace cv