用C++手撕NMS

用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;
}