NMS C++实现

转载https://blog.csdn.net/avideointerfaces/article/details/88551325

 

#include <iostream>
#include <vector>
using namespace std;

typedef struct{
	int x;
	int y;
	int width;
	int height;
}Rect;

typedef struct{
	Rect box;
	float score;
	int index;
}BBox;
bool compScore(BBox a, BBox b) {
	if(a.score>b.score) 
		return true;
	else 
		return false;
}
static float get_boxs_iou(Rect rect1, Rect rect2)
{
	int xx1, yy1, xx2, yy2;
 
	xx1 = max(rect1.x, rect2.x);
	yy1 = max(rect1.y, rect2.y);
	xx2 = min(rect1.x + rect1.width - 1, rect2.x + rect2.width - 1);
	yy2 = min(rect1.y + rect1.height - 1, rect2.y + rect2.height - 1);
 
	int insection_width, insection_height;
	insection_width = max(0, xx2 - xx1 + 1);
	insection_height = max(0, yy2 - yy1 + 1);
 
	float insection_area, union_area, iou;
	insection_area = float(insection_width) * insection_height;
	union_area = float(rect1.width*rect1.height + rect2.width*rect2.height - insection_area);
	iou = insection_area / union_area;
	return iou;
}
 
void nms(vector<Rect>&boxes,vector<float>&scores,float confThreshold, float nmsThreshold, vector<int> &indices ){
	BBox b_;
	vector<BBox> b_s;
	int i,j;

	//b_s包括 [box坐标,box score,box index]，保留index是为了下面根据前两项获取当前index
	for(i=0;i<boxes.size();i++){
		b_.box=boxes[i];
		b_.score=scores[i];
		b_.index=i;
		b_s.push_back(b_);
	}
	//从大到小排序
	sort(b_s.begin(),b_s.end(),compScore);

	int b_s_size=b_s.size();
	//对每一个
	for(i=0;i<b_s_size;i++){
		//如果当前box的iou值小于置信度，那么说明之后的都小于置信度，则退出
		if(b_s[i].score<confThreshold) 
			continue;
		//把当前box的index保留
		indices.push_back(i);
		//计算与当前box后面box 与 当前box 的iou，若iou较大则删除
		for(j=i+1;j<b_s_size;j++){
			float iou=get_boxs_iou(b_s[i].box,b_s[j].box);
			if(iou>nmsThreshold){
				b_s.erase(b_s.begin()+j);
				b_s_size=b_s.size();
			}
		}
	}
}