NMS和Soft-NMS的原理和Pytorch代碼實作

NMS都不會，做什麼Detection！

Non-maximum suppression（非極大值抑制）算法

NMS原理：

1. 首先得出所有的預測框集合

   B

   、 對應框的得分

   Scores

   ， NMS(IoU)門檻值

   T

   ;
2. 定義存放侯選框的集合

   H

   （初始為

   Null

   ）， 對

   Scores

   排序選出得分最大的框為

   maxBox

   ，

   将

   maxBox

   從集合

   B

   中移到集合H中，集合

   B

   中沒有

   maxBox

   框了;
3. 計算

   maxBox

   和

   B

   中剩餘的所有框的IoU, 将IoU大于

   T

   的從

   B

   中删除(認為和

   maxBox

   重疊了);
4. 重複2～3步驟，直到集合

   B

   為

   Null

   ， 集合H中存放的框就是NMS處理的結果；

   重複步驟是：

   （1）對集合B中剩餘框對應的得分進行排序， 選出最大得分的框maxBox，并從集合B中移到集合H中。

   （2） 計算這個得分最大的框maxBox和集合B中框的IoU門檻值，将大于IoU門檻值的框從B中删除。

NMS代碼實作

1. Pytorch代碼實作

from torch import Tensor
import torch


def box_area(boxes: Tensor) -> Tensor:
    """
    Computes the area of a set of bounding boxes, which are specified by its
    (x1, y1, x2, y2) coordinates.

    Arguments:
        boxes (Tensor[N, 4]): boxes for which the area will be computed. They
            are expected to be in (x1, y1, x2, y2) format

    Returns:
        area (Tensor[N]): area for each box
    """
    return (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])


def box_iou(boxes1: Tensor, boxes2: Tensor) -> Tensor:
    """
    Return intersection-over-union (Jaccard index) of boxes.

    Both sets of boxes are expected to be in (x1, y1, x2, y2) format.

    Arguments:
        boxes1 (Tensor[N, 4])
        boxes2 (Tensor[M, 4])

    Returns:
        iou (Tensor[N, M]): the NxM matrix containing the pairwise IoU values for every element in boxes1 and boxes2
    """
    area1 = box_area(boxes1)  # 每個框的面積 (N,)
    area2 = box_area(boxes2)  # (M,)

    lt = torch.max(boxes1[:, None, :2], boxes2[:, :2])  # [N,M,2] # N中一個和M個比較； 是以由N，M 個
    rb = torch.min(boxes1[:, None, 2:], boxes2[:, 2:])  # [N,M,2]

    wh = (rb - lt).clamp(min=0)  # [N,M,2]  #小于0的為0  clamp 鉗；夾鉗；
    inter = wh[:, :, 0] * wh[:, :, 1]  # [N,M]  

    iou = inter / (area1[:, None] + area2 - inter)
    return iou  # NxM， boxes1中每個框和boxes2中每個框的IoU值；


def nms(boxes: Tensor, scores: Tensor, iou_threshold: float):
    """
    :param boxes: [N, 4]， 此處傳進來的框，是經過篩選（NMS之前選取過得分TopK）之後， 在傳入之前處理好的；
    :param scores: [N]
    :param iou_threshold: 0.7
    :return:
    """
    keep = []  # 最終保留的結果， 在boxes中對應的索引；
    idxs = scores.argsort()  # 值從小到大的 索引

    while idxs.numel() > 0:  # 循環直到null； numel()： 數組元素個數
        # 得分最大框對應的索引, 以及對應的坐标
        max_score_index = idxs[-1]
        max_score_box = boxes[max_score_index][None, :]  # [1, 4]
        keep.append(max_score_index)
        if idxs.size(0) == 1:  # 就剩餘一個框了；
            break
        idxs = idxs[:-1]  # 将得分最大框 從索引中删除； 剩餘索引對應的框 和 得分最大框 計算IoU；
        other_boxes = boxes[idxs]  # [?, 4]
        ious = box_iou(max_score_box, other_boxes)  # 一個框和其餘框比較 1XM
        idxs = idxs[ious[0] <= iou_threshold]

    keep = idxs.new(keep)  # Tensor
    return keep
           

2. Pytorch代碼實作

import torch
def nms(boxes, scores, overlap=0.7, top_k=200):
    """
    輸入:
        boxes: 存儲一個圖檔的所有預測框。[num_positive,4].
        scores:置信度。如果為多分類則需要将nms函數套在一個循環内。[num_positive].
        overlap: nms抑制時iou的門檻值.
        top_k: 先選取置信度前top_k個框再進行nms.
    傳回:
        nms後剩餘預測框的索引.
    """

    keep = scores.new(scores.size(0)).zero_().long()
    # 儲存留下來的box的索引 [num_positive]
    # 函數new(): 建構一個有相同資料類型的tensor

    # 如果輸入box為空則傳回空Tensor
    if boxes.numel() == 0:
        return keep

    x1 = boxes[:, 0]  # x1 坐标
    y1 = boxes[:, 1]
    x2 = boxes[:, 2]
    y2 = boxes[:, 3]
    area = torch.mul(x2 - x1, y2 - y1)  # 并行化計算所有框的面積
    v, idx = scores.sort(0)  # 升序排序
    idx = idx[-top_k:]  # 前top-k的索引，從小到大
    xx1 = boxes.new()
    yy1 = boxes.new()
    xx2 = boxes.new()  # new() 無參數，建立 相同類型的空值；
    yy2 = boxes.new()
    w = boxes.new()
    h = boxes.new()

    count = 0
    while idx.numel() > 0:
        i = idx[-1]  # 目前最大score對應的索引  # 選取得分最大的框索引；
        keep[count] = i  # 存儲在keep中
        count += 1
        if idx.size(0) == 1:  # 跳出循環條件：box被篩選完了
            break
        idx = idx[:-1]  # 去掉最後一個

        # 剩下boxes的資訊存儲在xx，yy中
        torch.index_select(x1, 0, idx, out=xx1)  # 從x1中再次元0選取索引為idx 資料 輸出到xx1中；
        torch.index_select(y1, 0, idx, out=yy1)  # torch.index_select（） # 從tensor中按指定次元和索引 取值；
        torch.index_select(x2, 0, idx, out=xx2)
        torch.index_select(y2, 0, idx, out=yy2)

        # 計算目前最大置信框與其他剩餘框的交集，不知道clamp的同學确實容易被誤導
        xx1 = torch.clamp(xx1, min=x1[i])  # max(x1,xx1)  # x1 y1 的最大值
        yy1 = torch.clamp(yy1, min=y1[i])  # max(y1,yy1)
        xx2 = torch.clamp(xx2, max=x2[i])  # min(x2,xx2)  # x2 x3 最小值；
        yy2 = torch.clamp(yy2, max=y2[i])  # min(y2,yy2)
        w.resize_as_(xx2)
        h.resize_as_(yy2)
        w = xx2 - xx1  # w=min(x2,xx2)−max(x1,xx1)
        h = yy2 - yy1  # h=min(y2,yy2)−max(y1,yy1)
        w = torch.clamp(w, min=0.0)  # max(w,0)
        h = torch.clamp(h, min=0.0)  # max(h,0)
        inter = w * h

        # 計算目前最大置信框與其他剩餘框的IOU
        # IoU = i / (area(a) + area(b) - i)
        rem_areas = torch.index_select(area, 0, idx)  # 剩餘的框的面積
        union = rem_areas + area[i] - inter  # 并集
        IoU = inter / union  # 計算iou

        # 選出IoU <= overlap的boxes(注意le函數的使用)
        idx = idx[IoU.le(overlap)]  # le: 小于等于 傳回的bool ， 去除大于overlap的值；
    return keep, count
           

參考自：連結

3. Numpy代碼實作

import numpy as np
from numpy import array

def box_area(boxes :array):
    """
    :param boxes: [N, 4]
    :return: [N]
    """
    return (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])

def box_iou(box1 :array, box2: array):
    """
    :param box1: [N, 4]
    :param box2: [M, 4]
    :return: [N, M]
    """
    area1 = box_area(box1)  # N
    area2 = box_area(box2)  # M
    # broadcasting, 兩個數組各次元大小 從後往前對比一緻， 或者 有一次元值為1；
    lt = np.maximum(box1[:, np.newaxis, :2], box2[:, :2])
    rb = np.minimum(box1[:, np.newaxis, 2:], box2[:, 2:])
    wh = rb - lt
    wh = np.maximum(0, wh) # [N, M, 2]
    inter = wh[:, :, 0] * wh[:, :, 1]
    iou = inter / (area1[:, np.newaxis] + area2 - inter)
    return iou  # NxM

def numpy_nms(boxes :array, scores :array, iou_threshold :float):

    idxs = scores.argsort()  # 按分數 降序排列的索引 [N]
    keep = []
    while idxs.size > 0:  # 統計數組中元素的個數
        max_score_index = idxs[-1]
        max_score_box = boxes[max_score_index][None, :]
        keep.append(max_score_index)

        if idxs.size == 1:
            break
        idxs = idxs[:-1]  # 将得分最大框 從索引中删除； 剩餘索引對應的框 和 得分最大框 計算IoU；
        other_boxes = boxes[idxs]  # [?, 4]
        ious = box_iou(max_score_box, other_boxes)  # 一個框和其餘框比較 1XM
        idxs = idxs[ious[0] <= iou_threshold]

    keep = np.array(keep)  # Tensor
    return keep
           

Soft NMS原理和代碼

Soft-NMS的原理

原理步驟：

1. 首先得出所有的預測框集合

   B

   、 對應框的得分

   Scores

   ， NMS(IoU)門檻值

   T

   ;
2. 定義存放侯選框的集合

   H

   （初始為

   Null

   ），對

   Scores

   排序選出得分最大的框為

   maxBox

   ， 将

   maxBox

   從集合

   B

   中移到集合H中，集合

   B

   中沒有

   maxBox

   框了;
3. 計算

   maxBox

   和

   B

   中剩餘的所有框的IoU, 将IoU大于

   T

   的框的得分 按某種方式 降低（不删除了），
4. 重複2～3步驟，直到集合

   B

   為

   Null

   ， 集合H中存放的框就是Soft-NMS處理的結果；

   重複步驟是：

   （1）對集合B中剩餘框對應的得分進行排序(

   因為分數變化，必須排序

   )， 選出最大得分的框maxBox，并從集合B中移到集合H中。

   （2） 計算這個得分最大的框maxBox和集合B中框的IoU門檻值，将大于IoU門檻值的框對應的得分降低。

5. Soft-NMS傳回的結果是 框以及框對應的得分（得分是Soft-NMS抑制後的），說白了，就是抑制了框對應的得分， 使用時需要一個得分門檻值。

   

Soft-NMS權重函數的形式

M M M表示得分最大框， b i b_i bi​是除去得分最大框後剩餘的每個框。

原來的NMS可以描述如下：将IoU大于門檻值的視窗的得分全部置為0。

s i = { s i , i o u ( M , b i ) < T 0 , i o u ( M , b i ) > = T s_i = \left\{ \begin{array}{lr}s_i, iou(M,b_i) < T \\ 0, iou(M,b_i) >= T \end{array} \right. si​={si​,iou(M,bi​)<T0,iou(M,bi​)>=T​

(1) 線性權重抑制得分

s i = { s i , i o u ( M , b i ) < T s i ( 1 − i o u ( M , b i ) ) , i o u ( M , b i ) > = T s_i = \left\{ \begin{array}{lr}s_i, iou(M,b_i) < T \\ s_i(1 - iou(M, b_i)), iou(M,b_i) > =T \end{array} \right. si​={si​,iou(M,bi​)<Tsi​(1−iou(M,bi​)),iou(M,bi​)>=T​

(2) 高斯權重抑制得分

s i = s i e − i o u ( M , b i ) 2 σ , ∀ b i ∉ D s_i = s_ie^{-\frac{iou(M, b_i)^2}{\sigma}}, \forall b_i \notin D si​=si​e−σiou(M,bi​)2​,∀bi​∈/​D

NMS-Soft 實作代碼

Pytorch 代碼 未驗證 沒有合适比較代碼【歡迎指正錯誤】

from torch import Tensor
import torch


def box_area(boxes: Tensor) -> Tensor:
    return (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])


def box_iou(boxes1: Tensor, boxes2: Tensor) -> Tensor:
    area1 = box_area(boxes1)  # 每個框的面積 (N,)
    area2 = box_area(boxes2)  # (M,)

    lt = torch.max(boxes1[:, None, :2], boxes2[:, :2])  # [N,M,2] # N中一個和M個比較； 是以由N，M 個
    rb = torch.min(boxes1[:, None, 2:], boxes2[:, 2:])  # [N,M,2]

    wh = (rb - lt).clamp(min=0)  # [N,M,2]  # 删除面積小于0 不相交的  clamp 鉗；夾鉗；
    inter = wh[:, :, 0] * wh[:, :, 1]  # [N,M]  # 切片的用法 相乘次元減1
    iou = inter / (area1[:, None] + area2 - inter)
    return iou  # NxM， boxes1中每個框和boxes2中每個框的IoU值；


def soft_nms(boxes: Tensor, scores: Tensor, soft_threshold=0.01,  iou_threshold=0.7, weight_method=2, sigma=0.5):
    """
    :param boxes: [N, 4]， 此處傳進來的框，是經過篩選（選取的得分TopK）之後的
    :param scores: [N]
    :param iou_threshold: 0.7
    :param soft_threshold soft nms 過濾掉得分太低的框 （手動設定）
    :param weight_method 權重方法 1. 線性 2. 高斯
    :return:
    """
    keep = []
    idxs = scores.argsort()
    while idxs.numel() > 0:  # 循環直到null； numel()： 數組元素個數
        # 由于scores得分會改變，是以每次都要重新排序，擷取得分最大值
        idxs = scores.argsort()  # 評分排序
        if idxs.size(0) == 1:  # 就剩餘一個框了；
            keep.append(idxs[-1])  
            break
        keep_len = len(keep)
        max_score_index = idxs[-(keep_len + 1)]
        max_score_box = boxes[max_score_index][None, :]  # [1, 4]
        idxs = idxs[:-(keep_len + 1)]
        other_boxes = boxes[idxs]  # [?, 4]
        keep.append(max_score_index)  # 位置不能邊
        ious = box_iou(max_score_box, other_boxes)  # 一個框和其餘框比較 1XM
        # Soft NMS 處理， 和 得分最大框 IOU大于門檻值的框， 進行得分抑制
        if weight_method == 1:   # 線性抑制  # 整個過程 隻修改分數
            ge_threshod_bool = ious[0] >= iou_threshold
            ge_threshod_idxs = idxs[ge_threshod_bool]
            scores[ge_threshod_idxs] *= (1. - ious[0][ge_threshod_bool])  # 小于IoU門檻值的不變
            # idxs = idxs[scores[idxs] >= soft_threshold]  # 小于soft_threshold删除， 經過抑制後 門檻值會越來越小；
        elif weight_method == 2:  # 高斯抑制， 不管大不大于門檻值，都計算權重
            scores[idxs] *= torch.exp(-(ious[0] * ious[0]) / sigma) # 權重(0, 1]
            # idxs = idxs[scores[idxs] >= soft_threshold]
        # else:  # NMS
        #     idxs = idxs[ious[0] <= iou_threshold]

    # keep = scores[scores > soft_threshold].int()
    keep = idxs.new(keep)  # Tensor
    keep = keep[scores[keep] > soft_threshold]  # 最後處理門檻值
    boxes = boxes[keep]  # 保留下來的框
    scores = scores[keep]  # soft nms抑制後得分
    return boxes, scores
           

對應代碼和測試在NMS_SoftNMS項目，歡迎指正！