c++版的NMS(非极大抑制)实现

NMS(非极大抑制)是深度学习目标检测中常用的小算法,用来过滤掉同一个物体上的那些置信度较低的bbboxes,最后只剩下该目标检测框集中最大置信度的那个。

算法原理

   说它是小算法的原因是其原理很简单。

  1)先对输入检测框按置信度由高到低排序

2)挑选第一个检测框(即最高置信度,记为A)和其它检测框(记为B)进行iou计算

3)如果iou大于nmsThreshold, 那就将B清除掉

4)跳转到2)从剩余得框集里面找置信度最大得框和其它框分别计算iou

5)直到所有框都过滤完。

代码实现

网上得代码主要以python或matlab为主,我这里是c++实现得,接口参数和opencv里面得NMSBoxes(...)完全一样。代码里面有相应注释,而且经过测试,能正常工作得。

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
void nms_boxes(vector &boxes, vector &confidences, float confThreshold, float nmsThreshold, vector &indices)
{	
	BBOX bbox;
	vector bboxes;
	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(), comp);

	int updated_size = bboxes.size();
	for (i = 0; i < updated_size; i++)
	{
		if (bboxes[i].confidence < confThreshold)
			continue;
		indices.push_back(bboxes[i].index);
		for (j = i + 1; j < updated_size; j++)
		{
			float iou = get_iou_value(bboxes[i].box, bboxes[j].box);
			if (iou > nmsThreshold)
			{
				bboxes.erase(bboxes.begin() + j);
				updated_size = bboxes.size();
			}
		}
	}

小结

   做的比较匆忙,没太考虑效率。如果有更好得方法或有欠考虑周到得地方,恳请指正。

你可能感兴趣的:(深度学习)