您的位置:

NMS非极大值抑制

一、什么是NMS

NMS是非极大值抑制(Non-maximum Suppression)的缩写。它是目标检测领域非常重要的一个操作,用于抑制多个重复的候选框,确保最终选择的框是最具代表性的。

在目标检测中,通常会通过预测出一系列可能包含目标的候选框。但由于图像中的目标可能具有多个不同的尺寸、形状和方向,以及可能存在多个候选框的重叠,这些候选框可能会出现重复的情况。

因此,需要使用NMS来从重复的候选框中筛选出最最具代表性的框,以提高目标检测的精度。

二、NMS的实现原理

NMS的实现原理非常简单,它主要分为以下几个步骤:

1. 候选框排序

首先,将所有候选框按照置信度(confidence)进行排序。一般来说,置信度越高的框越有可能是目标框。

def sort_boxes(boxes, scores):
    """
    对所有的候选框按照置信度从大到小排序
    :param boxes: 所有候选框的坐标 (N, 4)
    :param scores: 所有候选框的置信度 (N,)
    :return: 排序后的候选框的下标
    """
    index = np.argsort(scores)[::-1]
    return index

2. 计算重叠面积

对于排序后的候选框,计算每两个框之间的IoU(Intersection of Union,即交集比)。

def iou(box1, box2):
    """
    计算两个框的IoU值
    :param box1: 第一个框的坐标 (x1, y1, x2, y2)
    :param box2: 第二个框的坐标 (x1, y1, x2, y2)
    :return: 两个框的IoU值
    """
    x1 = max(box1[0], box2[0])
    y1 = max(box1[1], box2[1])
    x2 = min(box1[2], box2[2])
    y2 = min(box1[3], box2[3])
    if x1 >= x2 or y1 >= y2:
        return 0
    intersection = (x2 - x1) * (y2 - y1)
    area1 = (box1[2] - box1[0]) * (box1[3] - box1[1])
    area2 = (box2[2] - box2[0]) * (box2[3] - box2[1])
    iou = intersection / (area1 + area2 - intersection)
    return iou

3. 进行抑制

对于排序后的每个框,依次与其后面的框进行比较,如果当前框与后面的框的IoU大于一定的阈值(如0.5),则将后面的框删除。

def nms(boxes, scores, iou_threshold=0.5):
    """
    NMS操作
    :param boxes: 所有候选框的坐标 (N, 4)
    :param scores: 所有候选框的置信度 (N,)
    :param iou_threshold: IoU阈值
    :return: 选择的候选框的下标
    """
    index = sort_boxes(boxes, scores)
    select_index = []
    while len(index) > 0:
        current_index = index[0]
        select_index.append(current_index)
        delete_index = []
        for i in range(1, len(index)):
            if iou(boxes[current_index], boxes[index[i]]) > iou_threshold:
                delete_index.append(i)
        index = np.delete(index, delete_index)
    return select_index

三、NMS的应用

NMS广泛应用于目标检测中,特别是基于深度学习的目标检测算法中。例如,Faster R-CNN、YOLO和SSD等算法中都使用了NMS。

以下是使用Python的OpenCV库实现NMS的完整示例:

import cv2
import numpy as np

def sort_boxes(boxes, scores):
    index = np.argsort(scores)[::-1]
    return index

def iou(box1, box2):
    x1 = max(box1[0], box2[0])
    y1 = max(box1[1], box2[1])
    x2 = min(box1[2], box2[2])
    y2 = min(box1[3], box2[3])
    if x1 >= x2 or y1 >= y2:
        return 0
    intersection = (x2 - x1) * (y2 - y1)
    area1 = (box1[2] - box1[0]) * (box1[3] - box1[1])
    area2 = (box2[2] - box2[0]) * (box2[3] - box2[1])
    iou = intersection / (area1 + area2 - intersection)
    return iou

def nms(boxes, scores, iou_threshold=0.5):
    index = sort_boxes(boxes, scores)
    select_index = []
    while len(index) > 0:
        current_index = index[0]
        select_index.append(current_index)
        delete_index = []
        for i in range(1, len(index)):
            if iou(boxes[current_index], boxes[index[i]]) > iou_threshold:
                delete_index.append(i)
        index = np.delete(index, delete_index)
    return select_index

if __name__ == '__main__':
    # example boxes and scores
    boxes = np.array([[0, 0, 50, 50], [30, 30, 80, 80], [60, 60, 100, 100], [80, 80, 120, 120]])
    scores = np.array([0.9, 0.8, 0.7, 0.6])

    # nms
    select_index = nms(boxes, scores, iou_threshold=0.5)

    # draw boxes
    img = np.zeros((200, 200, 3), dtype=np.uint8)
    for i in select_index:
        x1, y1, x2, y2 = boxes[i]
        cv2.rectangle(img, (x1, y1), (x2, y2), (0, 255, 0), 2)
    cv2.imshow('Image', img)
    cv2.waitKey()

运行上述代码后,将会显示如下图所示的输出结果:

四、总结

本文简要介绍了NMS非极大值抑制的基本原理和应用,以及使用Python的OpenCV库实现NMS的完整示例代码。对于深度学习领域中目标检测的相关研究与应用都有一定的参考意义。