非極大值抑制(nms)算法詳解[python]

一、起源

目標檢測在使用了基於深度學習的端到端模型後效果斐然。目前，常用的目標檢測算法，無論是One-stage的SSD系列算法、YOLO系列算法還是Two-stage的基於RCNN系列的算法，非極大值抑制都是其中必不可少的一個組件。在現有的基於anchor的目標檢測算法中，都會產生數量巨大的候選矩形框，這些矩形框有很多是指向同一目標，因此就存在大量冗餘的候選矩形框。非極大值抑制算法的目的正在於此，它可以消除多餘的框，找到最佳的物體檢測位置。

非極大值抑制（Non-Maximum Suppression，以下簡稱NMS算法）的思想是搜索局部極大值，抑制非極大值元素。針對不同的應用場景和檢測算法，由於矩形框的表徵方式不同，NMS算法具有各種變體，本文針對NMS及其各種變體（主要是本文檢測相關）進行介紹。

二、經典NMS

經典NMS最初第一次應用到目標檢測中是在RCNN算法中，其實現嚴格按照搜索局部極大值，抑制非極大值元素的思想來實現的，具體的實現步驟如下：

設定目標框的置信度閾值，常用的閾值是0.5左右
根據置信度降序排列候選框列表
選取置信度最高的框A添加到輸出列表，並將其從候選框列表中刪除
計算A與候選框列表中的所有框的IoU值，刪除大於閾值的候選框
重複上述過程，直到候選框列表爲空，返回輸出列表

其中IoU(Intersection over Union)爲交併比，如圖1所示，IoU相當於兩個區域交叉的部分除以兩個區域的並集部分得出的結果。圖2是IoU爲各個取值時的情況展示，一般來說，這個score ＞ 0.5 就可以被認爲一個不錯的結果了。

圖1

圖2

下面，通過一個具體例子來說明經典NMS究竟做了什麼。圖3的左圖是包含一個檢測目標（王鬧海）的實例圖片。其中的綠色矩形框代表了經過目標檢測算法後，生成的大量的帶置信度的Bounding box，矩形框左下角的浮點數即代表該Bounding box的置信度。在這裏，使用Python對經典NMS算法實現，並應用到該實例中去。當NMS的閾值設爲0.2時，最後的效果如圖3中右圖所示。

圖3

Python代碼如下：

本文所有代碼和數據的github地址：https://github.com/lzneu/letGo/tree/master/nms

def nms(bounding_boxes, Nt):
    if len(bounding_boxes) == 0:
        return [], []
    bboxes = np.array(bounding_boxes)

    # 計算 n 個候選框的面積大小
    x1 = bboxes[:, 0]
    y1 = bboxes[:, 1]
    x2 = bboxes[:, 2]
    y2 = bboxes[:, 3]
    scores = bboxes[:, 4]
    areas = (x2 - x1 + 1) * (y2 - y1 + 1)

    # 對置信度進行排序, 獲取排序後的下標序號, argsort 默認從小到大排序
    order = np.argsort(scores)

    picked_boxes = []  # 返回值
    while order.size > 0:
        # 將當前置信度最大的框加入返回值列表中
        index = order[-1]
        picked_boxes.append(bounding_boxes[index])

        # 獲取當前置信度最大的候選框與其他任意候選框的相交面積
        x11 = np.maximum(x1[index], x1[order[:-1]])
        y11 = np.maximum(y1[index], y1[order[:-1]])
        x22 = np.minimum(x2[index], x2[order[:-1]])
        y22 = np.minimum(y2[index], y2[order[:-1]])
        w = np.maximum(0.0, x22 - x11 + 1)
        h = np.maximum(0.0, y22 - y11 + 1)
        intersection = w * h

        # 利用相交的面積和兩個框自身的面積計算框的交併比, 將交併比大於閾值的框刪除
        ious = intersection / (areas[index] + areas[order[:-1]] - intersection)
        left = np.where(ious < Nt)
        order = order[left]
    return picked_boxes

三、soft-NMS

經典NMS是爲了去除重複的預測框，這種算法在圖片中只有單個物體被檢測的情況下具有很好的效果，如上圖所示，“腰子”王鬧海被成功的檢測出來。然而，經典NMS算法存在着一些問題：對於重疊物體無法很好的檢測。當圖像中存在兩個重疊度很高的物體時，經典NMS會過濾掉其中置信度較低的一個。如下圖所示，圖中存在兩個王鬧海，經典NMS過濾後的結果如下圖4所示：

圖4

而我們期望的結果是兩個腰子都被算法成功檢測出來。

爲了解決這類問題，Jan Hosang,等人提出了Soft-NMS算法。Soft-NMS的算法僞代碼如圖5所示。其中紅色框爲經典NMS的步驟，而綠色框中的內容爲Soft-NMS改進的步驟。可以看出，相對於經典NMS算法，Soft-NMS僅僅修改了一行代碼。當選取了最大置信度的Bounding box之後，計算其餘每個Bounding box與Bounding box的I ou值，經典NMS算法的做法是直接刪除Iou大於閾值的Bounding box；而Soft-NMS則是使用一個基於Iou的衰減函數，降低Iou大於閾值Nt的Bounding box的置信度，IoU越大，衰減程度越大。

換一個角度來理解，經典NMS同樣可以用一個基於Iou的衰減函數來表示，只是這個衰減函數是個0-1函數，當Iou大於閾值Nt，直接降低該Bounding box的置信度到0。關於經典NMS和Soft-NMS的衰減函數區別，可以通過如下的公式來表示[1]:

經典NMS：
Soft-NMS-線性：
Soft-NMS-高斯：

圖5

接下來，繼續使用圖4的實例來驗證Soft-NMS的效果，使用Python對Soft-NMS算法的實現如下

Python代碼：

def soft_nms(bboxes, Nt=0.3, sigma2=0.5, score_thresh=0.3, method=2):
    # 在 bboxes 之後添加對於的下標[0, 1, 2...], 最終 bboxes 的 shape 爲 [n, 5], 前四個爲座標, 後一個爲下標
    res_bboxes = deepcopy(bboxes)
    N = bboxes.shape[0]  # 總的 box 的數量
    indexes = np.array([np.arange(N)])  # 下標: 0, 1, 2, ..., n-1
    bboxes = np.concatenate((bboxes, indexes.T), axis=1)  # concatenate 之後, bboxes 的操作不會對外部變量產生影響
    # 計算每個 box 的面積
    x1 = bboxes[:, 0]
    y1 = bboxes[:, 1]
    x2 = bboxes[:, 2]
    y2 = bboxes[:, 3]
    scores = bboxes[:, 4]
    areas = (x2 - x1 + 1) * (y2 - y1 + 1)

    for i in range(N):
        # 找出 i 後面的最大 score 及其下標
        pos = i + 1
        if i != N - 1:
            maxscore = np.max(scores[pos:], axis=0)
            maxpos = np.argmax(scores[pos:], axis=0)
        else:
            maxscore = scores[-1]
            maxpos = 0
        # 如果當前 i 的得分小於後面的最大 score, 則與之交換, 確保 i 上的 score 最大
        if scores[i] < maxscore:
            bboxes[[i, maxpos + i + 1]] = bboxes[[maxpos + i + 1, i]]
            scores[[i, maxpos + i + 1]] = scores[[maxpos + i + 1, i]]
            areas[[i, maxpos + i + 1]] = areas[[maxpos + i + 1, i]]
        # IoU calculate
        xx1 = np.maximum(bboxes[i, 0], bboxes[pos:, 0])
        yy1 = np.maximum(bboxes[i, 1], bboxes[pos:, 1])
        xx2 = np.minimum(bboxes[i, 2], bboxes[pos:, 2])
        yy2 = np.minimum(bboxes[i, 3], bboxes[pos:, 3])
        w = np.maximum(0.0, xx2 - xx1 + 1)
        h = np.maximum(0.0, yy2 - yy1 + 1)
        intersection = w * h
        iou = intersection / (areas[i] + areas[pos:] - intersection)
        # Three methods: 1.linear 2.gaussian 3.original NMS
        if method == 1:  # linear
            weight = np.ones(iou.shape)
            weight[iou > Nt] = weight[iou > Nt] - iou[iou > Nt]
        elif method == 2:  # gaussian
            weight = np.exp(-(iou * iou) / sigma2)
        else:  # original NMS
            weight = np.ones(iou.shape)
            weight[iou > Nt] = 0
        scores[pos:] = weight * scores[pos:]
    # select the boxes and keep the corresponding indexes
    inds = bboxes[:, 5][scores > score_thresh]
    keep = inds.astype(int)
    return res_bboxes[keep]

圖6

四、Locality-Aware NMS

以上的NMS或相關算法均爲針對目標檢測問題而提出的，而文本檢測作爲目標檢測的一個特殊領域，具有其特殊之處，比如文本檢測中面臨的成千上萬的候選幾何體，如果使用計算複雜度爲的經典NMS類算法，由於n的數量過大，這是不可接受的。爲此曠世公司於2017在EAST（一種非常好用的One-stage的文本檢測算法，敬請關注交易只能團隊下期文章！）論文[2]中提出了Locality-Aware NMS算法。

Locality-Aware NMS算法考慮了文本檢測任務中臨近像素高度相關的特徵，提出了一種基於行的幾何體合併方法，它的步驟主要是兩個：

合併同一行區域中的幾何體積
對合並後的Bounding box集合進行標準的NMS操作

Locality-Aware NMS提升了算法的時間複雜度，在最優情況下可以提升至。即使是最壞情況，也不會比經典NMS的計算量更大。

Python代碼如下：

import numpy as np
from shapely.geometry import Polygon

def intersection(g, p):
    # 取g,p中的幾何體信息組成多邊形
    g = Polygon(g[:8].reshape((4, 2)))
    p = Polygon(p[:8].reshape((4, 2)))
    # 判斷g,p是否爲有效的多邊形幾何體
    if not g.is_valid or not p.is_valid:
        return 0
    # 取兩個幾何體的交集和並集
    inter = Polygon(g).intersection(Polygon(p)).area
    union = g.area + p.area - inter
    if union == 0:
        return 0
    else:
        return inter / union

def weighted_merge(g, p):
    # 取g,p兩個幾何體的加權（權重根據對應的檢測得分計算得到）
    g[:8] = (g[8] * g[:8] + p[8] * p[:8]) / (g[8] + p[8])
    # 合併後的幾何體的得分爲兩個幾何體得分的總和
    g[8] = (g[8] + p[8])
    return g

def standard_nms(S, thres):
    # 標準NMS
    order = np.argsort(S[:, 8])[::-1]
    keep = []
    while order.size > 0:
        i = order[0]
        keep.append(i)
        ovr = np.array([intersection(S[i], S[t]) for t in order[1:]])
        inds = np.where(ovr <= thres)[0]
        order = order[inds + 1]
    return S[keep]


def nms_locality(polys, thres=0.3):
    '''
    locality aware nms of EAST
    :param polys: a N*9 numpy array. first 8 coordinates, then prob
    :return: boxes after nms
    '''
    S = []  # 合併後的幾何體集合
    p = None  # 合併後的幾何體
    for g in polys:
        if p is not None and intersection(g, p) > thres:  # 若兩個幾何體的相交面積大於指定的閾值，則進行合併
            p = weighted_merge(g, p)
        else:  # 反之，則保留當前的幾何體
            if p is not None:
                S.append(p)
            p = g
    if p is not None:
        S.append(p)
    if len(S) == 0:
        return np.array([])
    return standard_nms(np.array(S), thres)

五、Softer-NMS

以上的NMS算法，無論是經典NMS還是Soft-NMS算法，都是在一種假設前提下：置信度最高的Bounding box就是目標的候選位置最精確的物體位置。但是事實上，這個假設可能並不成立，或者說並不那麼精確。針對這個問題，來自卡內基梅隆大學與曠視科技的研究人員在文中提出了一種新的非極大抑制算法Softer-NMS，顯著改進了目標檢測的定位精度，代碼已經開源，目前Github上的Star已超100，可謂短短兩天已經引起了不小的關注。作者關注了兩種目前NMS會出現的問題：