NMS最大值抑制的c++实现

本人是将pytorch训练后的模型转换为pt格式，然后利用libtorch调用，在进行推理后，得到tensor，需要进行后处理，要使用到NMS，所以查了一些资料，并完善了一下代码。

在.h文件中声明相关结构体与函数，需要的库头文件自行添加。

typedef struct {
		Rect box;
		float score;
		int index;
	}BBox;

void nms(vector<Rect> &boxes, vector<float> &scores, float confThreshold,
		float nmsThreshold, vector<int> &indices, Mat img);
float get_boxs_iou(Rect rect1, Rect rect2);

函数实现：

float CMFCLibtorchTestDlg::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 CMFCLibtorchTestDlg::nms(vector<Rect> &boxes, vector<float> &scores, float confThreshold, 
	float nmsThreshold, vector<int> &indices,Mat img) 
{
	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);
	sort(b_s.begin(), b_s.end(), [](const BBox &a, const BBox &b) { return a.score > b.score; });
	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);
                //用来绘制检测到的目标框
		indices.push_back(b_s[i].index);
		//rectangle(img, b_s[i].box, Scalar(0, 0, 255), 1, LINE_AA);
		//计算与当前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();
				j--;
			}
		}
	}
}

踩过的坑：在最后计算box和box的iou之后，大于阈值的则删除，此时，需要注意，删除后，vector所有的元素都会往前移动一个位置。因此，需要 j-- 操作。

if (iou > nmsThreshold) 
{			
	b_s.erase(b_s.begin() + j);
	b_s_size = b_s.size();
	j--;
}

终于通过单步调试，发现了该问题，菜是原罪啊。