NMS（非极大值抑制）实现_nmsboxes

在目标检测算法中，为了尽量不漏掉物体，会输出大量的检测结果（每一条结果由检测概率与检测框组成）。这些检测结果很可能有重复的，即会有多个框标出了同一个物体。
我们需要一个算法来过滤多余的检测框。最常用的算法就是NMS(Non-Maximum Suppresion, 非极大值抑制)。该算法的思路很简单：只保留局部概率最大的检测框，与其重合的其他检测框都会被舍去。

算法的伪代码如下：

输入所有检测概率、检测框
当还有检测框没有处理：
    从剩下的框里挑出检测概率最大的框 bbox_a
    遍历所有没有处理的框bbox_i：
        如果 bbox_i != bbox_a 且 bbox_i 与 bbox_a 重合：
            舍去 bbox_i
    把 bbox_a 输出成一个检测结果
1
2
3
4
5
6
7

当然，这个算法的描述还不够准确：究竟该怎么定义两个检测框是“重合”呢？如果两个检测框有交集就说它们重合是行不通的，因为图片中很可能会有挨得很近的物体，它们的检测框就是相交的。因此，我们一般用IoU（交并比）来描述两个框的重合程度，如果IoU超出某个阈值，就说这两个框是“重合”的。

IoU的计算很简单，算出两个矩形的「交面积」，再用两个矩形面积之和减去「交面积」就可以得到「并面积」，「交面积」比「并面积」就是IoU。

python实现：

import numpy as np


def NMSBoxes(input, score_threshold, nms_threshold):
    boxes = input[input[:, 4] > score_threshold] # score_threshold过滤

    x1 = boxes[:, 0] - boxes[:, 2] * 0.5
    y1 = boxes[:, 1] - boxes[:, 3] * 0.5
    x2 = boxes[:, 0] + boxes[:, 2] * 0.5
    y2 = boxes[:, 1] + boxes[:, 3] * 0.5

    scores = boxes[:, 4]
    areas = (x2-x1) * (y2-y1)
    results = []

    index = np.argsort(scores)[::-1]     # 按置信度进行排序

    while(index.size):
        # 置信度最高的框
        i = index[0]
        results.append(index[0])

        if(index.size == 1): # 如果只剩一个框，直接返回
            break

        # 计算交集左下角与右上角坐标
        inter_x1 = np.maximum(x1[i], x1[index[1:]])
        inter_y1 = np.maximum(y1[i], y1[index[1:]])
        inter_x2 = np.minimum(x2[i], x2[index[1:]])
        inter_y2 = np.minimum(y2[i], y2[index[1:]])

        # 计算交集的面积
        w = np.maximum(inter_x2 - inter_x1, 0) 
        h = np.maximum(inter_y2 - inter_y1, 0)
        inter_area = w * h
        
        # 计算当前框与其余框的iou
        iou = inter_area / (areas[index[1:]] + areas[i] - inter_area)
        id = np.where(iou < nms_threshold)[0]
        index = index[id + 1]

    return input[results]


if __name__ == '__main__':
    boxes = np.array([[1,1,3,3,0.6], [0,0,2,2,0.7], [0,0,2,2,0.2]])
    res = NMSBoxes(input=boxes, score_threshold=0.5, nms_threshold=0.2)
    print(res)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48

C++实现：

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


typedef struct Bbox
{
	float x;
	float y;
	float w;
	float h;
	float score;
}Bbox;


void printboxes(std::vector<Bbox> boxes)
{
	for (size_t i = 0; i < boxes.size(); i++)
	{
		printf("%f %f %f %f %f \n", boxes[i].x, boxes[i].y, boxes[i].w, boxes[i].h, boxes[i].score);
	}
	printf("\n");
}


float iou(Bbox box1, Bbox box2)
{
	float x1 = std::max(box1.x - box1.w * 0.5, box2.x - box2.w * 0.5);
	float y1 = std::max(box1.y - box1.h * 0.5, box2.y - box2.h * 0.5);
	float x2 = std::min(box1.x + box1.w * 0.5, box2.x + box2.w * 0.5);
	float y2 = std::min(box1.y + box1.h * 0.5, box2.y + box2.h * 0.5);
	float w = std::max(0.0f, x2 - x1);
	float h = std::max(0.0f, y2 - y1);
	float inter_area = w * h;
	float iou = inter_area / (box1.w * box1.h + box2.w * box2.h - inter_area);

	return iou;
}


std::vector<Bbox> NMSBoxes(std::vector<Bbox> & boxes, float score_threshold, float nms_threshold)
{
	std::vector<Bbox> boxes_results;
	std::sort(boxes.begin(), boxes.end(), [](Bbox box1, Bbox box2) {return box1.score > box2.score; }); //将box按照score从高到低排序

	//score_threshold过滤
	for (size_t i = 0; i < boxes.size(); i++)
	{
		if (boxes[i].score < score_threshold)
			boxes.erase(boxes.begin() + i);
		else
			i++;
	}
	//printboxes(boxes);

	while (boxes.size() > 0)
	{
		boxes_results.push_back(boxes[0]);
		int index = 1;
		//计算最大score对应的box与剩下所有的box的IOU,移除所有大于IOU阈值的box
		while (index < boxes.size())
		{
			float iou_value = iou(boxes[0], boxes[index]);
			//std::cout << iou_value << std::endl;

			if (iou_value > nms_threshold)
				boxes.erase(boxes.begin() + index);
			else
				index++;
		}
		boxes.erase(boxes.begin()); //results已经保存，所以可以将最大的删除了
	}

	return boxes_results;
}


int main()
{
	std::vector<Bbox> boxes = { { 1,1,3,3,0.6 }, { 0,0,2,2,0.7 }, { 0,0,2,2,0.2 } };
	std::vector<Bbox> res = NMSBoxes(boxes, 0.5, 0.2);
	printboxes(res);

	return 0;
}
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85

下面是在杜佬的代码中学到的实现方式，更容易理解，但需要多开辟vector的大小用来标记box是否要去除：

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


typedef struct Bbox
{
	float x;
	float y;
	float w;
	float h;
	float score;
}Bbox;


void printboxes(std::vector<Bbox> boxes)
{
	for (size_t i = 0; i < boxes.size(); i++)
	{
		printf("%f %f %f %f %f \n", boxes[i].x, boxes[i].y, boxes[i].w, boxes[i].h, boxes[i].score);
	}
	printf("\n");
}


float iou(Bbox box1, Bbox box2)
{
	float x1 = std::max(box1.x - box1.w * 0.5, box2.x - box2.w * 0.5);
	float y1 = std::max(box1.y - box1.h * 0.5, box2.y - box2.h * 0.5);
	float x2 = std::min(box1.x + box1.w * 0.5, box2.x + box2.w * 0.5);
	float y2 = std::min(box1.y + box1.h * 0.5, box2.y + box2.h * 0.5);
	float w = std::max(0.0f, x2 - x1);
	float h = std::max(0.0f, y2 - y1);
	float inter_area = w * h;
	float iou = inter_area / (box1.w * box1.h + box2.w * box2.h - inter_area);

	return iou;
}


std::vector<Bbox> NMSBoxes(std::vector<Bbox> & boxes, float score_threshold, float nms_threshold)
{
	std::vector<Bbox> boxes_results;
	std::sort(boxes.begin(), boxes.end(), [](Bbox box1, Bbox box2) {return box1.score > box2.score; }); //将box按照score从高到低排序

	//score_threshold过滤
	for (size_t i = 0; i < boxes.size(); i++)
	{
		if (boxes[i].score < score_threshold)
			boxes.erase(boxes.begin() + i);
		else
			i++;
	}
	//printboxes(boxes);

	std::vector<bool> remove_flags(boxes.size());
	for (size_t i = 0; i < boxes.size(); ++i) 
	{
		if (remove_flags[i]) 
			continue;

		boxes_results.push_back(boxes[i]);
		for (size_t j = i + 1; j < boxes.size(); ++j) 
		{
			if (remove_flags[j]) 
				continue;

			if (iou(boxes[i], boxes[j]) >= nms_threshold)
				remove_flags[j] = true;
		}
	}

	return boxes_results;
}


int main()
{
	std::vector<Bbox> boxes = { { 1,1,3,3,0.6 }, { 0,0,2,2,0.7 }, { 0,0,2,2,0.2 } };
	std::vector<Bbox> res = NMSBoxes(boxes, 0.5, 0.2);
	printboxes(res);

	return 0;
}
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84