c++版的NMS（非極大抑制）實現

轉載：c++版的NMS（非極大抑制）實現

一.演算法原理

1）先對輸入檢測框按置信度由高到低排序

2）挑選第一個檢測框(即最高置信度，記為A）和其它檢測框（記為B）進行iou計算

3）如果iou大於nmsThreshold， 那就將B清除掉

4）跳轉到2）從剩餘得框集裡面找置信度最大得框和其它框分別計算iou

5）直到所有框都過濾完。

二.程式碼實現

• 汗顏，幾乎照搬

typedef struct {
	Rect box;
	float confidence;
	int index;
}BBOX;

static float get_iou_value(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;
}

//input:  boxes: 原始檢測框集合;
//input:  confidences：原始檢測框對應的置信度值集合
//input:  confThreshold 和 nmsThreshold 分別是 檢測框置信度閾值以及做nms時的閾值
//output:  indices  經過上面兩個閾值過濾後剩下的檢測框的index
bool cmp(BBOX &b1, BBOX &b2) {
	return b1.confidence > b2.confidence;
}
vector<BBOX> nms_boxes(vector<Rect> &boxes, vector<float> &confidences, float confThreshold, float nmsThreshold, vector<int> &indices)
{
	BBOX bbox;
	vector<BBOX> bboxes,ans;
	vector<bool> need_delete(boxes.size(), false);
	int i, j;
	for (i = 0; i < boxes.size(); i++)
	{
		bbox.box = boxes[i];
		bbox.confidence = confidences[i];
		bbox.index = i;
		bboxes.push_back(bbox);
	}
	sort(bboxes.begin(), bboxes.end(), cmp);

	for (i = 0; i < bboxes.size(); i++)
	{
		if (need_delete[i])
			continue;
		if (bboxes[i].confidence < confThreshold)
			break;
		indices.push_back(bboxes[i].index);
		for (j = i + 1; j < bboxes.size(); j++)
		{
			if (need_delete[j])
				continue;
			float iou = get_iou_value(bboxes[i].box, bboxes[j].box);
			if (iou > nmsThreshold)
			{
				need_delete[j] = true;

			}
		}
		return ans;
	}