NMS(非极大值抑制)算法非极大值抑制是用于减少物体检测算法中重叠边界框或区域的数量的技术。通过对每个类别的检测框按置信度排序,然后逐个遍历,保留置信度最高的框,并抑制与其重叠且置信度低的框,从而得到更准确和简洁的检测结果。
假设我们使用一个人脸检测算法来识别一张图像中的多个人脸。该算法可能会在图像中产生多个边界框,每个边界框对应一个被检测到的人脸。然而,算法可能会产生多个重叠的边界框,每个边界框都对应同一个人脸。
在这种情况下,我们可以应用非极大值抑制来减少重叠边界框的数量,以确保最终的检测结果准确而简洁。具体步骤如下:
1、对于每个检测到的人脸边界框,计算其置信度(例如,检测得分)。
2、根据置信度对所有边界框进行排序,将置信度最高的边界框作为起始点。
3、从置信度第二高的边界框开始,逐个遍历每个边界框。
4、对于当前正在考虑的边界框,计算其与前面已选定的边界框的重叠程度(例如,使用交并比IoU)。
5、如果当前边界框与已选定的边界框重叠程度高于预设的阈值,则将其抑制(即不保留);否则保留当前边界框。
6、重复上述步骤,直到所有边界框都被处理完毕。
通过这个过程,最终每个人脸只会被保留一个最具代表性的边界框,从而避免了多个重叠边界框的情况,得到了准确而简洁的人脸检测结果。
draw_bbox.py
在图像上绘制边界框及其置信度
- import cv2
-
- def draw_box(img, bbox, confidence=None, offset=(0, 0)):
- """
- 绘制图像上的一个边界框
- :param img: 待绘制的图像
- :param bbox: 边界框坐标,必须是(xmin, ymin, xmax, ymax)的格式
- :param confidence: 边界框的置信度
- :param offset: 偏移量,用于调整绘制的边界框位置
- :return: 绘制了边界框的图像
- """
- # 将bbox转换为整数类型
- x1, y1, x2, y2 = [int(i) for i in bbox]
-
- # 根据偏移量调整边界框的位置
- x1 += offset[0]
- x2 += offset[0]
- y1 += offset[1]
- y2 += offset[1]
-
- # 边界框的颜色和置信度文本
- color = 0.5 # 边界框颜色
- label = str(confidence) # 置信度文本
-
- # 获取文本大小
- t_size = cv2.getTextSize(label, cv2.FONT_HERSHEY_PLAIN, 1, 2)[0]
-
- # 绘制边界框
- cv2.rectangle(img, (x1, y1), (x2, y2), color, 3) # 绘制边界框
- cv2.rectangle(img, (x1, y1), (x1 + t_size[0] + 3, y1 + t_size[1] + 4), color, -1) # 绘制文本背景
- cv2.putText(img, label, (x1, y1 + t_size[1] + 4), cv2.FONT_HERSHEY_PLAIN, 2, [255, 255, 255], 2) # 绘制置信度文本
-
- return img
nms.py
实现了非极大值抑制(NMS)算法,用于在目标检测任务中去除重叠的边界框
- import numpy as np
- import cv2
-
- from draw_bbox import draw_box # 导入绘制边界框的函数
-
- def nms(bboxes, scores, iou_thresh):
- """
- 非极大值抑制(NMS)算法,用于去除重叠的边界框
- :param bboxes: 边界框列表,每个边界框格式为[xmin, ymin, xmax, ymax]
- :param scores: 边界框对应的置信度列表
- :param iou_thresh: IOU(交并比)阈值,用于判断两个边界框是否重叠
- :return: 经过NMS处理后的边界框和置信度列表
- """
- # 将边界框坐标拆分为四个坐标列表
- x1 = bboxes[:, 0]
- y1 = bboxes[:, 1]
- x2 = bboxes[:, 2]
- y2 = bboxes[:, 3]
- areas = (y2 - y1) * (x2 - x1) # 计算每个边界框的面积
-
- # 结果列表
- result = []
- index = scores.argsort()[::-1] # 对置信度进行从高到低的排序,并获取对应的索引
-
- # 循环处理边界框,直至列表为空
- while index.size > 0:
- i = index[0] # 取出当前置信度最高的边界框
- result.append(i) # 将当前边界框的索引加入结果列表
-
- # 计算当前边界框与其他边界框的IOU
- x11 = np.maximum(x1[i], x1[index[1:]])
- 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)
- h = np.maximum(0, y22 - y11 + 1)
- overlaps = w * h
- ious = overlaps / (areas[i] + areas[index[1:]] - overlaps)
-
- # 仅保留IOU低于阈值的边界框索引
- idx = np.where(ious <= iou_thresh)[0]
- index = index[idx + 1] # 更新剩余边界框的索引列表
-
- # 返回经过NMS处理后的边界框和置信度
- return bboxes[result], scores[result]
-
-
- if __name__ == '__main__':
- # 读取原始图像
- raw_img = cv2.imread('./test.png')
-
- # 原始的边界框和置信度
- bboxes = [[183, 625, 269, 865], [197, 603, 296, 853], [190, 579, 295, 864], [537, 507, 618, 713], [535, 523, 606, 687]]
- confidences = [0.7, 0.9, 0.95, 0.9, 0.6]
-
- # 绘制未经过NMS处理的原始检测结果
- img = raw_img.copy()
- for x, y in zip(bboxes, confidences):
- img = draw_box(img, x, y)
- cv2.imwrite("./org_test.png", img) # 保存结果图像
-
- # 进行NMS处理
- bboxes, scores = nms(np.array(bboxes), np.array(confidences), 0.5)
-
- # 绘制经过NMS处理后的检测结果
- img = raw_img.copy()
- for x, y in zip(list(bboxes), list(scores)):
- img = draw_box(img, x, y)
- cv2.imwrite("./nms_test.png", img) # 保存结果图像
测试图片
未使用nms算法
使用nms算法