【深度学习实战（29）】后处理之NMS（非极大值抑制）

一、NMS工作原理

NMS 的工作原理：

1. 置信度排序：对于每个类别，NMS 首先根据每个边界框的置信度（即预测框中含有目标的概率）进行排序。
2. 选择最高置信度框：从置信度最高的边界框开始，将其作为当前考虑的“最大”候选。
3. 计算交并比（IoU）：计算当前最大候选与所有其他边界框的交并比（Intersection over Union）。IoU 是两个边界框交集面积与并集面积的比值，用于衡量边界框之间的重叠程度。
4. 抑制重叠框：如果某个边界框与当前最大候选的 IoU 高于预设的阈值（例如，0.5），则认为它们检测到的是同一个目标，因此将该边界框从候选列表中移除。
5. 更新候选列表：移除所有被抑制的边界框后，从剩余的边界框中选择置信度最高的作为新的“最大”候选。
6. 迭代过程：重复步骤3-5，直到所有边界框都被处理完毕。

二、NMS C++代码实现

#include <vector>
#include <iostream>
#include <algorithm>

struct BBox 
{
	float x, y, w, h, conf;
	float getArea() const { return w* h; }
};


float iou(const BBox a, const BBox b)
{
	float area_a = a.getArea();
	float area_b = b.getArea();

	float x1 = std::max(a.x - a.w / 2, b.x - b.w / 2);
	float y1 = std::max(a.y - a.h / 2, b.y - b.h / 2);
	float x2 = std::min(a.x + a.w / 2, b.x + b.w / 2);
	float y2 = std::min(a.y + a.h / 2, b.y + b.h / 2);
	std::cout << "x1 = " << x1 << std::endl;
	std::cout << "y1 = " << y1 << std::endl;
	std::cout << "x2 = " << x2 << std::endl;
	std::cout << "y2 = " << y2 << std::endl;

	//有可能没有交集
	float area_jiao = std::max(0.f,(x2 - x1))*std::max(0.f,(y2 - y1));
	std::cout << "area_jiao = " << area_jiao << std::endl;


	//并集
	float area_bing = area_a + area_b - area_jiao;
	std::cout << "area_jiao / area_bing = " << area_jiao / area_bing << std::endl;

	float iou_ = area_jiao > 0 ? area_jiao / area_bing : 0;

	std::cout << "iou_ = " << iou_ << std::endl;


	return area_jiao > 0 ? area_jiao / area_bing : 0;

}


std::vector<BBox> nms(std::vector<BBox> &in_box, float thres)
{
	//按置信度排序
	std::sort(in_box.begin(), in_box.end(), [](const BBox & a, const BBox & b)
	{
		return a.conf < b.conf;
	});

	std::vector<BBox> results;
	for (size_t i = 0; i < in_box.size(); i++)
	{
		bool keep = true;
		for (size_t j = 0; j < results.size(); j++)
		{
			//如果当前BBox和results中BBox重叠度大于阈值，就抑制
			//否则，就放入results中
			if (iou(in_box[i], results[j]) > thres)
			{
				keep = false;
			}
			//抑制掉，就结束本次results遍历
			break;
		}

		if (keep)
		{
			results.push_back(in_box[i]);
		}
	}
	return results;
}

int main() {

	//设置一些BBox
	std::vector<BBox> boxes;
	boxes.push_back({ 100, 100, 50, 50, 0.9 });
	boxes.push_back({ 102, 101, 50, 50, 0.8 });
	boxes.push_back({ 98, 98, 50, 50, 0.7 });
	boxes.push_back({ 102, 98, 50, 50, 0.85 });

	//nms
	float thres = 0.82;
	std::vector<BBox> ret = nms(boxes, thres);

	//输出
	for (size_t i = 0; i < ret.size(); i++)
	{
		std::cout << "x=" << ret[i].x << "y=" << ret[i].y << "w=" << ret[i].w << "h=" << ret[i].h << "conf=" << ret[i].conf << std::endl;
	}

	system("pause");
	return 0;

}

run

三、soft NMS介绍

当物体密集时，传统NMS方法可能会错误地移除一些正确的检测框。而Soft-NMS考虑了所有可能的检测框，因此可以更准确地处理这种情况。soft-nms的核心就是降低置信度。比如一张人脸上有3个重叠的bounding box, 置信度分别为0.9, 0.7, 0.85 。选择得分最高的建议框，经过第一次处理过后，得分变成了0.9, 065, 0.55（此时将得分最高的保存在D中）。这时候再选择第二个bounding box作为得分最高的，处理后置信度分别为0.65, 0.45（这时候3个框也都还在），最后选择第三个，处理后得分不改变。最终经过soft-nms抑制后的三个框的置信度分别为0.9, 0.65, 0.45。最后设置阈值，将得分si小于阈值的去掉。