C++手写NMS

文章目录

  • 前言
  • 一、NMS是什么?
  • 二、代码展示
  • 三、代码实现思路
  • 总结


前言

目标检测模型推理后,一般都需要进行NMS操作进行多余框去重,板端部署一般不用opencv自带的NMS,所以记录下手写NMS的代码。


一、NMS是什么?

非极大值抑制(Non-Maximum Suppression,NMS)算法,用于在一组边界框(BoundingBoxes)中去除高度重叠的边界框,只保留最具代表性的一个。

二、代码展示

#include 
#include 
#include 

// 边界框结构
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;
}

三、代码实现思路

iou(const BoundingBox& a, const BoundingBox& b): 这个函数计算两个边界框a和b之间的交并比(IoU,Intersection over Union)。它首先计算每个边界框的面积,然后计算它们的交集区域的坐标范围和面积。最后,它将交集区域的面积除以并集区域的面积,得到IoU值,即重叠度。

nms(std::vector& boxes, float threshold): 这个函数实现了非极大值抑制算法。它接受一个包含边界框的向量 boxes 和一个阈值 threshold 作为参数。首先,它按照边界框的置信度(confidence)对边界框进行降序排序,以便首先处理置信度最高的边界框。

然后,它遍历每个边界框,并检查它是否应该被保留。对于每个边界框,它将其与已经被保留的边界框进行比较,计算它们之间的IoU。如果IoU大于阈值 threshold,则说明两个边界框高度重叠,当前边界框不被保留;否则,当前边界框被保留。

最终,函数返回一个向量 result,其中包含经过非极大值抑制后被保留下来的边界框。


总结

NMS通常用于目标检测中,以去除重叠的检测结果,只保留最具代表性的结果,以提高检测的准确性。
如果阅读本文对你有用,欢迎点赞收藏呀!!!
2023年9月27日15:58:10
C++手写NMS_第1张图片

你可能感兴趣的:(C++,深度学习,c++,开发语言,计算机视觉,NMS,后处理)