C++ handwritten NMS


Preface

After inferring the target detection model, it is generally necessary to perform NMS operations to remove redundant frames. Board-end deployment generally does not use the NMS that comes with opencv, so the handwritten NMS code is recorded.


1. What is NMS?

The Non-Maximum Suppression (NMS) algorithm is used to remove highly overlapping bounding boxes in a set of bounding boxes (BoundingBoxes), retaining only the most representative one.

2. Code display

#include <iostream>
#include <vector>
#include <algorithm>

// 边界框结构
struct BoundingBox {
    
    
    float x, y, w, h, confidence;

    // 计算边界框的面积
    float area() const {
    
    
        return w * h;
    }
};

// 计算两个边界框之间的IoU
float iou(const BoundingBox& a, const BoundingBox& b) {
    
    
    const float area_a = a.area();
    const float area_b = b.area();

    // 计算重叠区域的坐标范围
    const float x1 = std::max(a.x, b.x);
    const float y1 = std::max(a.y, b.y);
    const float x2 = std::min(a.x + a.w, b.x + b.w);
    const float y2 = std::min(a.y + a.h, b.y + b.h);

    // 计算重叠区域的面积
    const float intersection_area = std::max(0.0f, x2 - x1) * std::max(0.0f, y2 - y1);

    // 计算并集区域的面积
    const float union_area = area_a + area_b - intersection_area;

    // 计算IoU
    return union_area > 0 ? intersection_area / union_area : 0;
}

// 非极大值抑制函数
std::vector<BoundingBox> nms(std::vector<BoundingBox>& boxes, float threshold) {
    
    
    // 根据置信度排序
    std::sort(boxes.begin(), boxes.end(),
              [](const BoundingBox& a, const BoundingBox& b) {
    
     return a.confidence > b.confidence; });

    std::vector<BoundingBox> result;
    for (size_t i = 0; i < boxes.size(); ++i) {
    
    
        bool keep = true;
        for (size_t j = 0; j < result.size(); ++j) {
    
    
            if (iou(boxes[i], result[j]) > threshold) {
    
    
                keep = false;
                break;
            }
        }
        if (keep) {
    
    
            result.push_back(boxes[i]);
        }
    }
    return result;
}

int main() {
    
    
    std::vector<BoundingBox> boxes;

    // 添加一些示例边界框数据
    boxes.push_back({
    
    10, 10, 20, 20, 0.9});
    boxes.push_back({
    
    15, 15, 25, 25, 0.8});
    boxes.push_back({
    
    30, 30, 20, 20, 0.7});
    boxes.push_back({
    
    40, 40, 15, 15, 0.85});

    // 运行非极大值抑制算法,保留不重叠的边界框
    float threshold = 0.5;
    std::vector<BoundingBox> result = nms(boxes, threshold);

    // 输出保留下来的边界框
    std::cout << "保留的边界框:" << std::endl;
    for (const BoundingBox& box : result) {
    
    
        std::cout << "x: " << box.x << ", y: " << box.y << ", w: " << box.w << ", h: " << box.h << ", confidence: " << box.confidence << std::endl;
    }

    return 0;
}

3. Code implementation ideas

iou(const BoundingBox& a, const BoundingBox& b): This function calculates the intersection and union ratio between two bounding boxes a and b (IoU,Intersection over Union). It first calculates the area of ​​each bounding box, and then calculates the coordinate range and area of ​​their intersection region. Finally, it divides the area of ​​the intersection region by the area of ​​the union region to obtain the IoU value, which is the degree of overlap.

nms(std::vector<BoundingBox>& boxes, float threshold): This function implements the non-maximum suppression algorithm. It accepts a vector containing bounding boxes boxes and a threshold as arguments. (confidence)First, it sorts the bounding boxes in descending order by their confidence so that the bounding boxes with the highest confidence are processed first.

Then it goes through each bounding box and checks if it should be preserved. For each bounding box, it compares it to the already preserved bounding box and calculates the IoU between them. If the IoU is greater than the threshold, it means that the two bounding boxes are highly overlapping and the current bounding box is not retained; otherwise, the current bounding box is retained.

Finally, the function returns a vector result containing the bounding boxes that have been preserved after non-maximum suppression.


Summarize

NMS is often used in target detection to remove overlapping detection results and retain only the most representative results to improve detection accuracy.
If reading this article is useful to you, please like and save it! ! !
September 27, 2023 15:58:10
Insert image description here

Guess you like

Origin blog.csdn.net/JulyLi2019/article/details/133353948