C++笔记--基于C++实现NMS算法

1--完整代码

// g++ nms.cpp -std=c++11 -o nms_test
// ./nms_test

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

// 锚框结构体
struct BoundingBox{
    float x, y, width, height, score;
};

class NMS{
public:
    NMS(){};
    // 非极大值抑制函数
    std::vector<BoundingBox> nonMaximumSuppression(const std::vector<BoundingBox>& boxes, float overlapThreshold){
        std::vector<BoundingBox> selectedBoxes; // 返回的锚框

        // 根据分数降序排序
        std::vector<BoundingBox> sortedBoxes = boxes;
        std::sort(sortedBoxes.begin(), sortedBoxes.end(), [](const BoundingBox& a, const BoundingBox& b){
            return a.score > b.score;
        });

        while(!sortedBoxes.empty()){
            // 每次取分数最高的合法锚框
            const BoundingBox& currentBox = sortedBoxes.front();
            selectedBoxes.push_back(currentBox);
            sortedBoxes.erase(sortedBoxes.begin()); // 取完后从候选锚框中删除

            // 计算剩余锚框与合法锚框的IOU
            std::vector<BoundingBox>::iterator it = sortedBoxes.begin();
            while (it != sortedBoxes.end()){
                const BoundingBox& candidateBox = *it; // 取当前候选锚框
                float intersection = intersectionArea(currentBox, candidateBox); // 计算候选框和合法框的交集面积
                float iou = intersection / (boxArea(currentBox) + boxArea(candidateBox) - intersection); // 计算iou
                if (iou >= overlapThreshold) it = sortedBoxes.erase(it);  // 根据阈值过滤锚框，过滤完it指向下一个锚框
                else it++; // 保留当前锚框，判断下一个锚框
            }
        }
        return selectedBoxes;
    }

    // 计算两个矩形框的交集面积
    float intersectionArea(const BoundingBox& box1, const BoundingBox& box2){
        float x1 = std::max(box1.x, box2.x);
        float y1 = std::max(box1.y, box2.y);
        float x2 = std::min(box1.x + box1.width, box2.x + box2.width);
        float y2 = std::min(box1.y + box1.height, box2.y + box2.height);
        if(x1 < x2 && y1 < y2) return (x2 - x1) * (y2 - y1);
        else return 0.0f;
    }

    // 计算矩形框的面积
    float boxArea(const BoundingBox& box){
        return box.width * box.height;
    }

};

int main(int argc, char *argv[]){
    // 生成一组锚框测试数据
    std::vector<BoundingBox> boxes = {{10, 10, 20, 20, 0.9},
                                      {15, 15, 25, 25, 0.85},
                                      {30, 30, 20, 20, 0.7},
                                      {12, 12, 22, 22, 0.95}};

    // 设定过滤阈值
    float overlapThreshold = 0.5;

    // nms过滤
    NMS nms;
    std::vector<BoundingBox> selectedBoxes = nms.nonMaximumSuppression(boxes, overlapThreshold);

    // 遍历打印过滤后的合法框
    for(const BoundingBox& box : selectedBoxes){
        std::cout << "Selected Box: (x=" << box.x << ", y=" << box.y << ", width=" << box.width 
                  << ", height=" << box.height << ", score=" << box.score << ")" << std::endl;
    }
    return 0;
}

2--编译运行

g++ nms.cpp -std=c++11 -o nms_test
./nms_test