NMS(非极大值抑制)实现

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

算法的伪代码如下:

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

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

IoU的计算很简单,算出两个矩形的「交面积」,再用两个矩形面积之和减去「交面积」就可以得到「并面积」,「交面积」比「并面积」就是IoU。
NMS(非极大值抑制)实现_第1张图片

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)

C++实现:

#include 
#include 
#include 


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


void printboxes(std::vector<Bbox> vec_boxes)
{
	for (size_t i = 0; i < vec_boxes.size(); i++)
	{
		printf("%f %f %f %f %f \n", vec_boxes[i].x, vec_boxes[i].y, vec_boxes[i].w, vec_boxes[i].h, vec_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> & vec_boxes, float score_threshold, float nms_threshold)
{
	std::vector<Bbox> results;
	std::sort(vec_boxes.begin(), vec_boxes.end(), [](Bbox box1, Bbox box2) {return box1.score > box2.score; }); //将box按照score从高到低排序

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

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

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

	return 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;
}

你可能感兴趣的:(#,object,detection,deep,learning,目标检测,NMS)