NMS的理解及Python實現

NMS的定義

全稱爲non maximum suppression，翻譯成中文是非極大抑制。怎麼理解呢？
這裏上一張圖。

如上圖所示，綠色的框識別出了鳴人，紅色的識別出了水門，藍色的則識別出了自來也。
但是，我們對於每個對象只需要一個框就夠了，在這裏，無論是鳴人，水門還是自來也都有三個框，我們應該保留哪個框呢？當然是得分（預測概率）最高的哪個。
NMS就是幹這個活的，也就是，不是得分最高的（極大值）的框都可以滾蛋了。

NMS實現

步驟1：按照得分的順序將所有bounding box進行排序。
步驟2：最高的那個肯定留下，然後計算它與其他所有bounding box（得分都低於他）的IoU（交併比，不解釋了）。若某個bounding box的ioU小於設定的閾值，那就留下他，不然就刪掉他。
步驟3：餘下的繼續執行1~2步
這裏舉個例子：
給定上圖的3個人物的3個bounding box及得分數組[x1, y1, x2, y2, score]。
第一行三個是鳴人，第二行三個是水門，第三行三個是自來也。

dets = [[218, 322, 385, 491, 0.98],[247, 312, 419, 461, 0.83],[237, 344, 407, 510, 0.92],
            [757, 218, 937, 394, 0.96],[768, 198, 962, 364, 0.85],[740, 240, 906, 414, 0.83],
            [1101, 84, 1302, 303, 0.82], [1110, 67, 1331, 260, 0.97], [1123, 42, 1362, 220, 0.85]]

我們按照上述步驟可以發現，得分最高的是鳴人的某個框(0.98)，那我們保留，並計算該框與其他所有框的IoU。這樣會形成8個IoU值（原本9個，拿出一個和其他8個比）。遍歷該IoU數組，假設閾值爲0.5，如果ioU小於我們的閾值，我們肯定留下它，否則就幹掉它。這裏注意，雖然需要和8個bounding box對比，但是實際上鳴人的框與自來也還有水門的框的IoU是0，所以一定會保留水門和自來也的框，也就是該輪操作至少剩下6個框。

overlap [0.         0.         0.62888921 0.         0.         0.
 0.55520685 0.        ]

結果我們看到，果然，我們要保留這6個0，而IoU爲0.62和0.55的那倆框其實就是鳴人的除最高得分的另外兩個框，怎麼辦？刪除就好了。
剩下的就是水門和自來也啦。

Python代碼實現

import numpy as np

def py_cpu_nms(dets, thres):
    x1 = dets[:, 0]
    y1 = dets[:, 1]
    x2 = dets[:, 2]
    y2 = dets[:, 3]
    scores = dets[:, 4]
    area = (x2 - x1 + 1) * (y2 - y1 + 1)
    order = scores.argsort()[::-1]
    print("order", order)

    # needs to save
    keep = []
    while order.size > 0:
        i = order[0]
        keep.append(i)
        xx1 = np.maximum(x1[i], x1[order[1:]])
        yy1 = np.maximum(y1[i], y1[order[1:]])
        xx2 = np.minimum(x2[i], x2[order[1:]])
        yy2 = np.minimum(y2[i], y2[order[1:]])

        w = np.maximum(0.0, xx2-xx1+1)
        h = np.maximum(0.0, yy2-yy1+1)

        inter = w * h
        overlap = inter / (area[i] + area[order[1:]] - inter)
        print("overlap", overlap)

        inds = np.where(overlap <= thres)[0]
        print("ind", inds)
        order = order[inds+1]
        print("order", order)

    return keep

if __name__ == "__main__":
    dets = [[218, 322, 385, 491, 0.98],[247, 312, 419, 461, 0.83],[237, 344, 407, 510, 0.92],
            [757, 218, 937, 394, 0.96],[768, 198, 962, 364, 0.85],[740, 240, 906, 414, 0.83],
            [1101, 84, 1302, 303, 0.82], [1110, 67, 1331, 260, 0.97], [1123, 42, 1362, 220, 0.85]]
    # [757, 218, 937, 394, 0.96],[768, 198, 962, 364, 0.85],[740, 240, 906, 414, 0.83]
    dets = np.array(dets)
    print(py_cpu_nms(dets, 0.5))

代碼運行過程如下：

order [0 7 3 2 8 4 5 1 6]
overlap [0.         0.         0.62888921 0.         0.         0.
 0.55520685 0.        ]
ind [0 1 3 4 5 7]
order [7 3 8 4 5 6]
overlap [0.         0.59778611 0.         0.         0.64035466]
ind [0 2 3]
order [3 4 5]
overlap [0.63086943 0.61164895]
ind []
order []
result [0, 7, 3]

按照我們剛纔輸入的順序，我們發現，最終留下的bounding box恰恰就是鳴人，水門和自來也得分最高的那三個。
代碼還是比較簡單，就不解釋啦。

加油，爲了未來！爲了瓜瓜！

NMS的理解及Python實現