用C++手撕NMS
NMS主要用于去除目标检测网络最后在目标周围大量重叠的bounding box,该部分用C++实现需要定义bounding box结构体,之后对bounding box列表进行排序,优先选择与目标框iou高的bounding box,与其他的bounding box比较iou的值,高于阈值的删去。
代码过程如下所示
#include<iostream>
#include<vector>
#include<algorithm>
using namespace std;
typedef struct bBox {
int x;
int y;
int w;
int h;
float score;
}bBox;
float iou(bBox box1, bBox box2) {
int x1 = box1.x > box2.x ? box1.x:box2.x;
int y1 = box1.y > box2.y ? box1.y:box2.y;
int x2 = (box1.x + box1.w) > (box2.x + box2.w) ? (box1.x + box1.w):(box2.x + box1.w);
int y2 = (box1.y + box1.h) > (box2.y + box2.h) ? (box1.y + box1.h):(box2.y + box1.h);
int w = (x2 - x1) > 0 ? (x2 - x1) : 0;
int h = (y2 - y1) > 0 ? (y2 - y1) : 0;
float over_area = w * h;
return over_area / (box1.w*box1.h + box2.w*box2.h - over_area);
}
//写一下排序的回调函数,得分高的优先
static bool sort_score(bBox box1,bBox box2) {
return box1.score > box2.score ? true : false;
}
vector<bBox> nms(std::vector<bBox>&boxes, float threshold)
{
vector<bBox>resluts;
std::sort(boxes.begin(), boxes.end(), sort_score);
while (boxes.size()> 0)
{
resluts.push_back(boxes[0]);
int index = 1;
while (index < boxes.size()) {
float iou_value = iou(boxes[0], boxes[index]);
if (iou_value > threshold) {
boxes.erase(boxes.begin() + index);
}
else {
index++;
}
}
boxes.erase(boxes.begin());
}
return resluts;
}
int main() {
vector<bBox> dataBox;
bBox box1 = { 9,10,12,6,0.9 };
bBox box2 = { 10,0,10,20,0.6 };
dataBox.push_back(box1);
dataBox.push_back(box2);
vector<bBox> res;
res = nms(dataBox, 0.6);
for (int i = 0; i < res.size(); i++) {
printf("%d %d %d %d %f", res[i].x, res[i].y, res[i].w, res[i].h, res[i].score);
cout << endl;
}
return 0;
}