[第九夜] 人臉偵測 : MTCNN 的實做 Part.2-2 -- 使用 NMS 對預測框去蕪存菁

15th鐵人賽

zivzhong

2023-09-24 23:59:23

517 瀏覽

分享至

前言

我們的 AI model 在讀入一張照片之後可能會預測出多個可能框出來，如下方示例：

這樣的結果展示出同一個人臉上看起來多了好多冗框（雖說確實還是比沒有框好啦），但我們更傾向下面這樣的結果：

那我們要怎麼做到這件事呢？仔細想想是不是就是

如果都在差不多的地方的框就可以結合並且去掉多餘的
那這也是我們今天要介紹的技術，透過 NMS 來去掉多餘框

NMS 介紹

非最大值抑制（Non-Maximum Suppression，簡稱NMS）是一種廣泛用於物體檢測和邊界框回歸的技術。在MTCNN（Multi-task Cascaded Convolutional Networks）中，NMS的主要作用是去除多個重疊的候選人臉框，僅保留最可能包含人臉的邊界框。

NMS的計算原理：

NMS的目標是根據物體檢測的信心分數（或概率）來選擇最佳的邊界框，並去除重疊度較高的邊界框。以下是NMS的計算流程：

Step.1 排序：將所有候選的邊界框按照其信心分數（概率）降序排列，即最高分數的邊界框排在前面。

Step.2 初始化選擇的框集合：創建一個空的已選擇框的集合。

Step.3 選擇最高分數框：從排序後的邊界框列表中選擇信心分數最高的邊界框（通常是第一個），將其添加到已選擇的框集合中。

Step.4 計算重疊度：計算剩餘邊界框與已選擇框集合中的框之間的IoU（交并比）。IoU是衡量兩個框重疊程度的指標，詳細計算可以參考下圖

Step.5 去除重疊框：去除與已選擇框集合中任何框的IoU大於某個閾值（通常是0.5）的邊界框。

Step.6 重複選擇：重複步驟3至步驟5，直到所有的邊界框都被處理完畢。

Step.7 返回選擇的框：最終已選擇的框集合中包含了不重疊且信心分數最高的邊界框。

NMS 實做

那因為我們已經開使用 Pytorch 在實做了，所以我們這邊也就盡量使用 Pytorch 來實做大家看看，下面是使用Python和PyTorch實現NMS的示例程式，假設已經有了候選框的信心分數（scores）和邊界框的坐標（boxes）：

import torch

def non_maximum_suppression(boxes, scores, threshold=0.5):
    """
    使用NMS選擇最佳的邊界框。
    
    參數：
    - boxes：(N, 4)的Tensor，包含N個候選框的坐標，每個框表示為(x1, y1, x2, y2)。
    - scores：(N,)的Tensor，包含N個候選框的信心分數。
    - threshold：IoU閾值，用於確定兩個框是否重疊。
    
    返回值：
    - selected_indices：已選擇的框的索引。
    """
    # 將框的信心分數按降序排列
    sorted_indices = torch.argsort(scores, descending=True)
    
    # 初始化已選擇的框的集合
    selected_indices = []
    
    while len(sorted_indices) > 0:
        # 選擇信心分數最高的框
        current_index = sorted_indices[0]
        selected_indices.append(current_index)
        
        # 計算選擇框與其他框的IoU
        current_box = boxes[current_index]
        other_boxes = boxes[sorted_indices[1:]]
        iou = calculate_iou(current_box, other_boxes)
        
        # 去除IoU大於閾值的框
        indices_to_remove = torch.where(iou > threshold)[0] + 1
        sorted_indices = torch.tensor(
            [idx for idx in sorted_indices if idx not in indices_to_remove])
    
    return selected_indices

def calculate_iou(box, boxes):
    """
    計算一個框與一組框之間的IoU。
    
    參數：
    - box：(4,)的Tensor，表示一個框的坐標。
    - boxes：(N, 4)的Tensor，表示N個框的坐標。
    
    返回值：
    - iou：(N,)的Tensor，表示box與每個框之間的IoU。
    """
    x1 = torch.maximum(box[0], boxes[:, 0])
    y1 = torch.maximum(box[1], boxes[:, 1])
    x2 = torch.minimum(box[2], boxes[:, 2])
    y2 = torch.minimum(box[3], boxes[:, 3])
    
    intersection = torch.maximum(0, x2 - x1) * torch.maximum(0, y2 - y1)
    area_box = (box[2] - box[0]) * (box[3] - box[1])
    area_boxes = (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])
    union = area_box + area_boxes - intersection
    
    iou = intersection / union
    return iou

以上這段程式實現了NMS(non_maximum_suppression)算法，non_maximum_suppression函數接受候選框的坐標和信心分數，然後根據IoU閾值選擇最佳的框。另外，還可以根據不同的 Task 需要調整IoU閾值（threshold）以適應特定應用場景，這裡只是先使用常見的 0.5 為例子。

NMS 可能的問題以及解法：SoftNMS

從上述算法中其實不難看出NMS有的一個潛在問題：

即嚴格的閾值可能導致過多的邊界框被刪除，就可能錯誤地排除了一些有用的信息。

我們舉下例子來說：

上面兩批馬各自會有一個框框，但綠框與紅框 overlap 的比例異常的高，那根據傳統 NMS 算法就可能刪掉綠框QQ

那 AI 圈後來有提出一個堅決的方式：Soft NMS，簡單說就是通過在選擇性保留和適度減少信心分數來解決這個問題（如下述 Step.5）。

以下是Soft NMS的計算流程：

Step.1 將所有候選的邊界框按照信心分數（概率）降序排列。

Step.2 初始化已選擇的框的集合為空。

Step.3 選擇信心分數最高的框，將其添加到已選擇的框集合中。

Step.4 計算該框與其他框的IoU。

Step.5 降低已選擇框和其他框的信心分數，以減少其競爭力。

Step.6 重複步驟3至步驟5，直到所有的邊界框都被處理完畢。

Step.7 返回已選擇的框集合。

以下是 NMS（藍色） v.s SoftNMS （紅色）的效果差異

上面可以看到長頸鹿們非常的靠近，左邊 NMS 就會刪掉一些長頸鹿的框框，而右邊的 softNMS 則盡可能的保留下來了！

Soft-NMS 實做

那既然與主要流程與NMS 差距不大只差在 step.5，那我們就直接看 code 吧～
以下是使用Python和PyTorch實現Soft NMS的程式：

import torch

def soft_non_maximum_suppression(boxes, scores, threshold=0.5, sigma=0.5):
    """
    使用Soft NMS選擇最佳的 bounding box。
    
    參數：
    - boxes：(N, 4)的Tensor，包含N個候選框的坐標，每個框表示為(x1, y1, x2, y2)。
    - scores：(N,)的Tensor，包含N個候選框的信心分數。
    - threshold：IoU閾值，用於確定兩個框是否重疊。
    - sigma：衰減參數，控制信心分數的衰減速度。
    
    返回值：
    - selected_indices：已選擇的框的索引。
    """
    # 將框的信心分數按降序排列
    sorted_indices = torch.argsort(scores, descending=True)
    
    # 初始化已選擇的框的集合
    selected_indices = []
    
    while len(sorted_indices) > 0:
        # 選擇信心分數最高的框
        current_index = sorted_indices[0]
        selected_indices.append(current_index)
        
        # 計算選擇框與其他框的IoU
        current_box = boxes[current_index]
        other_boxes = boxes[sorted_indices[1:]]
        iou = calculate_iou(current_box, other_boxes)
        
        # 計算信心分數的衰減因子
        decay = torch.exp(-(iou ** 2) / sigma)
        
        # 降低其他框的信心分數
        scores[sorted_indices[1:]] *= decay
        
        # 刪除信心分數過低的框
        indices_to_remove = torch.where(scores[sorted_indices[1:]] < threshold)[0] + 1
        sorted_indices = torch.tensor(
            [idx for idx in sorted_indices if idx not in indices_to_remove])
    
    return selected_indices

def calculate_iou(box, boxes):
    # 計算IoU的函數與前面的相同，這裡省略
    pass

# 使用示例
boxes = torch.tensor([[100, 100, 200, 200], [150, 150, 250, 250], [300, 300, 400, 400]])
scores = torch.tensor([0.9, 0.8, 0.7])
selected_indices = soft_non_maximum_suppression(boxes, scores, threshold=0.5, sigma=0.5)
print("選擇的框索引：", selected_indices)

以上的程式就是實現了Soft NMS，你可以複製下來試看看選擇最佳的邊界框並降低其競爭框的信心分數。可以更改 #使用示例 那一邊來看看各種情況的結果～～
值得一提的是 Soft NMS的一個重要參數是衰減參數（Decay Parameter），它控制了信心分數的衰減速度。較高的衰減參數會導致信心分數下降得較快，而較低的衰減參數則會使信心分數下降得較慢。這允許保留一些重疊的邊界框，同時降低其信心分數。

結語

今天我們介紹的 Object detection 中一個非常重要以及常見的計算--NMS，並且也介紹了他的進階版本，那之後的程式中我們也會使用這個 function。感謝大家今晚的參與，其他明天繼續與大家相見～～

參考資料

Bodla, Navaneeth, et al. "Soft-NMS--improving object detection with one line of code." Proceedings of the IEEE international conference on computer vision. 2017.