非极大值抑制算法（Non-Maximum Suppression，NMS）

Posted 2022-01-17 dlvector

算法思想

目标检测中常用到NMS，在faster R-CNN中，每一个bounding box都有一个打分，NMS实现逻辑是：

1，按打分最高到最低将BBox排序 ，例如：A B C D E F

2，A的分数最高，保留。从B-E与A分别求重叠率IoU，假设B、D与A的IoU大于阈值，那么B和D可以认为是重复标记去除

3，余下C E F，重复前面两步。

代码实现

"""
文件说明：https://www.cnblogs.com/king-lps/p/9031568.html
"""

# !/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Mon May  7 21:45:37 2018

@author: lps
"""
import numpy as np

# size(6,4)
boxes = np.array([[100, 100, 210, 210, 0.72],
                  [250, 250, 420, 420, 0.8],
                  [220, 220, 320, 330, 0.92],
                  [100, 100, 210, 210, 0.72],
                  [230, 240, 325, 330, 0.81],
                  [220, 230, 315, 340, 0.9]])


def py_cpu_nms(dets, thresh):
	# dets:(m,5)  thresh:scaler

	x1 = dets[:, 0]
	y1 = dets[:, 1]
	x2 = dets[:, 2]
	y2 = dets[:, 3]

	areas = (y2 - y1 + 1) * (x2 - x1 + 1)
	scores = dets[:, 4]
	keep = []

	index = scores.argsort()[::-1]

	while index.size > 0:
		i = index[0]  # every time the first is the biggst, and add it directly
		keep.append(i)

		x11 = np.maximum(x1[i], x1[index[1:]])  # calculate the points of overlap
		y11 = np.maximum(y1[i], y1[index[1:]])
		x22 = np.minimum(x2[i], x2[index[1:]])
		y22 = np.minimum(y2[i], y2[index[1:]])

		w = np.maximum(0, x22 - x11 + 1)  # the weights of overlap
		h = np.maximum(0, y22 - y11 + 1)  # the height of overlap

		overlaps = w * h

		ious = overlaps / (areas[i] + areas[index[1:]] - overlaps)

		idx = np.where(ious <= thresh)[0]

		index = index[idx + 1]  # because index start from 1

	return keep


import matplotlib.pyplot as plt


def plot_bbox(dets, c='k'):
	x1 = dets[:, 0]
	y1 = dets[:, 1]
	x2 = dets[:, 2]
	y2 = dets[:, 3]

	plt.plot([x1, x2], [y1, y1], c)  # c for color
	plt.plot([x1, x1], [y1, y2], c)
	plt.plot([x1, x2], [y2, y2], c)
	plt.plot([x2, x2], [y1, y2], c)
	plt.title("after nms")
	plt.show()


plot_bbox(boxes, 'k')  # before nms

keep = py_cpu_nms(boxes, thresh=0.7)
plot_bbox(boxes[keep], 'r')  # after nms

输出结果：
nms之前

nms之后

这里有几个细节点：

1. 获取下一轮的点坐标时，需要+1，这是因为iou计算是index[1:]，抛掉了概率最大值的那个框，所以少一个size.
   

2. 计算IOU时，请选用矩阵坐标系统来思考（即opencv的坐标系，左上角是原点，往下是ｘ递增的方向，往右是ｙ递增的方向），其实只要按照矩阵row和col的格式来思考就可以了．相交的左上角都是max,相交区域的右下角都是min．

参考

NMS的python实现
python numpy.where()函数的用法
非极大值抑制算法（Non-Maximum Suppression，NMS）
非极大值抑制（NMS）的几种实现