物体検出で重なったバウンディングボックスを除去・集約するアルゴリズムのまとめ (NMS, Soft-NMS, NMW, WBF)

物体検出の分野では、検出した物体をバウンディングボックス (BBox) で囲んで、それぞれに信頼度 (スコア) を算出します。

このとき重複したBBoxを除去あるいは集約するアルゴリズムにはバリエーションがあります。物体検出モデルの後処理やコンペなどでよく使われる4つを紹介します。

NMS
Soft-NMS
NMW
WBF

最初におさらい: IoU (Intersection over Union)

2つのBBoxがどれくらい重複しているかを表す指標の1つで、1.0に近づくほど重複しています。

分子が重なっている面積
分母が2つのBBoxの総面積

$IoU = \frac{Intersection}{Union}$

実装

def iou(a: tuple, b: tuple) -> float:
    a_x1, a_y1, a_x2, a_y2 = a
    b_x1, b_y1, b_x2, b_y2 = b
    
    if a == b:
        return 1.0
    elif (
        (a_x1 <= b_x1 and a_x2 > b_x1) or (a_x1 >= b_x1 and b_x2 > a_x1)
    ) and (
        (a_y1 <= b_y1 and a_y2 > b_y1) or (a_y1 >= b_y1 and b_y2 > a_y1)
    ):
        intersection = (min(a_x2, b_x2) - max(a_x1, b_x1)) * (min(a_y2, b_y2) - max(a_y1, b_y1))
        union = (a_x2 - a_x1) * (a_y2 - a_y1) + (b_x2 - b_x1) * (b_y2 - b_y1) - intersection
        return intersection / union
    else:
        return 0.0
    
assert iou((10, 20, 20, 30), (10, 20, 20, 30)) == 1.0
assert iou((10, 20, 20, 30), (20, 20, 30, 30)) == 0.0
assert 0.142 < iou((10, 20, 20, 30), (15, 25, 25, 35)) < 0.143
assert 0.111 < iou((10, 10, 40, 40), (20, 20, 30, 30)) < 0.112

NMS (Non-Maximum Suppression)

IoUがある閾値を超えて重なっているBBoxの集合から、スコアが最大のBBoxを残して、それ以外を除去するのがNMS¹です。

f:id:ohke:20200620094403p:plain

実装

def nms(bboxes: list, scores: list, iou_threshold: float) -> list:
    new_bboxes = []
    
    while len(bboxes) > 0:
        i = scores.index(max(scores))
        bbox = bboxes.pop(i)
        scores.pop(i)
        
        deletes = []
        for j, (bbox_j, score_j) in enumerate(zip(bboxes, scores)):
            if iou(bbox, bbox_j) > iou_threshold:
                deletes.append(j)
                
        for j in deletes[::-1]:
            bboxes.pop(j)
            scores.pop(j)
                
        new_bboxes.append(bbox)
        
    return new_bboxes
    
assert nms([(10, 20, 20, 30), (20, 20, 30, 30), (15, 25, 25, 35)], [0.8, 0.7, 0.9], 0.1) == [(15, 25, 25, 35)]

Soft-NMS

NMSではスコアが最大のBBox以外は除去されていましたが、例えば以下のように、同じラベルの物体が重なりあっている場合に誤って除去してしまう (再現率を落とす) ことが起こりえます。

f:id:ohke:20200620092744p:plain

こうしたケースに着目し、IoU閾値を超えたBBoxを残しつつ、代わりにスコアを割り引くのがSoft-NMS²です。

f:id:ohke:20200620133111p:plain

割引きの計算式は以下2つが提案されていますが、いずれもIoUが大きくなる (= より重複している) とより多く割り引くようになっています。

実装

先程のNMSの実装を少し変えるだけで実現できます。割り引き計算には上の1つ目の式を用いています。

def soft_nms(bboxes: list, scores: list, iou_threshold: float) -> (list, list):
    new_bboxes, new_scores = [], []
    
    while len(bboxes) > 0:
        i = scores.index(max(scores))
        bbox = bboxes.pop(i)
        score = scores.pop(i)
        
        for j, (bbox_j, score_j) in enumerate(zip(bboxes, scores)):
            iou_j = iou(bbox, bbox_j)
            if iou_j > iou_threshold:
                scores[j] = score_j * (1 - iou_j)
                
        new_bboxes.append(bbox)
        new_scores.append(score)
        
    return new_bboxes, new_scores
    
print(soft_nms([(10, 20, 20, 30), (20, 20, 30, 30), (15, 25, 25, 35)], [0.8, 0.7, 0.9], 0.1))
# ([(15, 25, 25, 35), (10, 20, 20, 30), (20, 20, 30, 30)], [0.9, 0.6857142857142858, 0.6])

Non-Maximum Weighted (NMW)

NMSとSoft-NMSは、いずれもBBoxの座標情報を変更していません。以下のように、いずれのBBoxでも正確に捉えきれていない場合、BBoxのどれかを残すのではなく、"良い感じに"ミックスして平均的なBBoxを作るほうが全体を上手に捉えることができそうな気がしてきます。

f:id:ohke:20200620231231p:plain

NMW³は、重なりあったBBoxをスコアとIoUで重み付けして足し合わせることで、1つの新たなBBoxを作り出すアルゴリズムです。

$b_{pre}$ が新たに作るBBoxです
$b_{argmax_{i} c_{i}}$ はスコアが最も高いBBoxです (つまりNMSなら残されるBBoxです)

実装

IoUが低いのでそれほど目立ってないですが、わずかにBBox (15, 25, 25, 35) が他の2つに近づいていることがわかります。

def nmw(bboxes: list, scores: list, iou_threshold: float) -> list:
    new_bboxes = []
    
    while len(bboxes) > 0:
        i = scores.index(max(scores))
        bbox = bboxes.pop(i)
        score = scores.pop(i)
        
        numerator = (bbox[0] * score, bbox[1] * score, bbox[2] * score, bbox[3] * score)
        denominator = score
        deletes = []
        for j, (bbox_j, score_j) in enumerate(zip(bboxes, scores)):
            iou_j = iou(bbox, bbox_j)
            if iou_j > iou_threshold:
                w = scores[j] * iou_j
                numerator = (
                    numerator[0] + bbox_j[0] * w, numerator[1] + bbox_j[1] * w,
                    numerator[2] + bbox_j[2] * w, numerator[3] + bbox_j[3] * w
                )
                denominator += w
                
                deletes.append(j)
                
        for j in deletes[::-1]:
            bboxes.pop(j)
            scores.pop(j)
          
        new_bboxes.append((
            numerator[0] / denominator, numerator[1] / denominator,
            numerator[2] / denominator, numerator[3] / denominator
        ))
        
    return new_bboxes
    
print(nmw([(10, 20, 20, 30), (20, 20, 30, 30), (15, 25, 25, 35)], [0.8, 0.7, 0.9], 0.1))
# [(14.935897435897434, 24.038461538461537, 24.935897435897434, 34.03846153846154)]

Weighted Boxes Fusion (WBF)

これまで触れてきたアルゴリズムは、重なり合わないBBoxについては何もせず残します。同じ画像に対して複数の推論を行う場合 (複数モデルでアンサンブルする、オーグメンテーションをかけるなど) 、これが誤検出につながるケースがあります。

3つのモデルで"dog"を検出する推論をかけ、以下のようにBBoxが出力されたとします。ある1つのモデルだけ左の人を検出してしまっています。NMS・Soft-NMS・NMWでは、この左の誤ったBBoxを除去したりスコアを下げたりできません。

WBF⁴は、複数モデルの推論結果を集約することを主目的としたアルゴリズムです。

複数モデルで推論させて得られたBBoxから、IoUを閾値として重複するBBoxを見つけ出します。このT個の重複したBBoxについて、スコア (C) と座標 (X1, X2, Y1, Y2) を以下の式で計算します。スコアで重み付けされた平均値となっており、この時点ではIoUを用いていません。

全てのBBoxについてスコアと座標の計算が終わった後、以下のいずれかの式でスコアをモデルの数 (N) で正規化します。これにより、検出されたモデルの数 (つまりT) が少ないBBoxほどスコアが下がるため、上のように1つのモデルだけで検出されたBBoxをスコアで足切りすることができるようになります。

実装

以下の通り、重なり合うBBoxは高いスコア (0.8) をキープしますが、重なっていないBBox (30, 30, 40, 40) は元のスコアの1/3されています。

def wbf(bboxes: list, scores: list, iou_threshold: float, n: int) -> (list, list):
    lists, fusions, confidences = [], [], []
    
    indexes = sorted(range(len(bboxes)), key=scores.__getitem__)[::-1]
    for i in indexes:
        new_fusion = True
        
        for j in range(len(fusions)):
            if iou(bboxes[i], fusions[j]) > iou_threshold:
                lists[j].append(bboxes[i])
                confidences[j].append(scores[i])
                fusions[j] = (
                    sum([l[0] * c for l, c in zip(lists[j], confidences[j])]) / sum(confidences[j]),
                    sum([l[1] * c for l, c in zip(lists[j], confidences[j])]) / sum(confidences[j]),
                    sum([l[2] * c for l, c in zip(lists[j], confidences[j])]) / sum(confidences[j]),
                    sum([l[3] * c for l, c in zip(lists[j], confidences[j])]) / sum(confidences[j]),
                )
                
                new_fusion = False
                break

        if new_fusion:
            lists.append([bboxes[i]])
            confidences.append([scores[i]])
            fusions.append(bboxes[i])
            
        print(lists, fusions, confidences)
            
    confidences = [(sum(c) / len(c)) * (min(n, len(c)) / n) for c in confidences]
    
    return fusions, confidences
    
print(wbf([(10, 10, 20, 20), (15, 10, 25, 20), (15, 15, 25, 25), (30, 30, 40, 40)], [0.8, 0.7, 0.9, 0.5], 0.1, 3))
# ([(13.333333333333332, 11.874999999999998, 23.33333333333333, 21.874999999999996), (30, 30, 40, 40)], [0.8000000000000002, 0.16666666666666666])

使い分け

ではこれらアルゴリズムをどう使い分けるべきかについてです。最終的にはメトリクスを見ながら適切なものを選ぶべきかと思いますが、大まかにはこういった使い分けになるかと思います。

単一モデルの精度を高めたい ->
- 同一物体に対する誤検出が多い -> NMS (あるいは、IoU閾値を下げる or スコア閾値を上げる)
- 同一画像内にたくさんの物体が含まれており、重複による未検出が多い -> Soft-NMS (あるいは、IoU閾値を上げる or スコア閾値を下げる)
- 変形や動きが大きい物体が含まれており、1つのBBoxで捉えきれずにIoUが低い -> NMW
精度がある程度高いモデルを複数組み合わせて、より高い精度を達成したい -> WBF

まとめ

物体検出で重なったBBoxを除去・集約するアルゴリズムについて整理しました。

[1311.2524] Rich feature hierarchies for accurate object detection and semantic segmentation (初出の論文は別かも…?)↩
[1704.04503] Soft-NMS -- Improving Object Detection With One Line of Code ↩
Huajun Zhou, Zechao Li, Chengcheng Ning, and Jinhui Tang. Cad: Scale invariant framework for real-time object detection. In Proceedings of the IEEE International Conference on Computer Vision, pages 760–768, 2017.↩
[1910.13302] Weighted Boxes Fusion: ensembling boxes for object detection models ↩