pythonでSoft-NMSを実装する(1)アルゴリズム実装例

pythonでSoft-NMS(Soft Non-Maximum Suppression)を実装する方法を紹介します。Soft-NMSは、SSDやYOLOといった物体検出AIの後処理で使用されるNMS(Non-Maximum Suppression)の改良型アルゴリズムです。

この記事で紹介する実装例を高速化したサンプルコードを次の記事で紹介していますので参考にしてみてください。

この記事でできること

ニューラルネットワーク(ディープラーニング)等を用いた機械学習による物体検出AIでは、たくさんの検出枠(この記事では矩形と呼んでいます)が出力されます。

従来はNMSという処理で、重複する矩形の統合や削除の処理を行いますが、この記事ではNMSの改良型であるSoft-NMSの実装例を紹介します。

次の画像は、いらすとやの「離れて給食を食べる子供のイラスト」の二人の子供の顔付近にいくつかの矩形を付与した画像です。

この画像にSoft-NMSを適用した例が次の画像です。

ただし、Soft-NMSは矩形を直接削除する処理ではなく、重複する矩形の信頼度(スコア)を下げるアルゴリズムで、その部分がNMSと大きく違います。

この画像は、Soft-NMSを適用した後、信頼度が小さい矩形を削除した後の矩形を描画したものです。閾値を設定して、閾値以下の矩形を消去しました。

Soft-NMS(Soft Non-Maximum Suppression)とは

従来のNMSは物体検出AIで出力されるたくさんの重複した矩形（長方形）をマージしてすっきりさせるアルゴリズムですが、Soft-NMSでは矩形をマージし削除するのではなく、重複している矩形の信頼度（score）を下げるアルゴリズムとなっています。

信頼度（score）が閾値以下の矩形を削除する処理をSoft-NMSの後段に入れることで、NMSよりも高精度にscoreの低い矩形の棄却を行います。他にも物体検出の評価指標であるmAPの向上に寄与するなど、通常のNMSよりも検出精度に有利な働きをすることが多いようです。

Soft-NMSは次の論文で提案されました。

Bodla, N., Singh, B., Chellappa, R., Davis, L.S.: Soft-nms—improving object detection with one line of code. In: 2017 IEEE International Conference on Computer Vision (ICCV). pp. 5562–5570. IEEE (2017).
https://arxiv.org/abs/1704.04503

従来NMSの実装例や処理内容については、以前の記事で紹介しましたので参考までにリンクを貼っておきます。

numpyを使ったNMSの高速化版は次の記事で紹介しました。

Soft-NMS(Soft Non-Maximum Suppression)のアルゴリズム

Soft-NMSの計算手順は、具体的に次のような処理となります。

1. scoreが最大の矩形iを選ぶ
2. scoreがiより小さい全ての矩形jとのIoUをそれぞれ計算する
3. 関数f(score, IoU)で矩形jのscoreを全て更新する
4. 矩形iを出力（1のscore最大の矩形選択の対象から外す）
5. 1に戻る（未出力の矩形の中でscore最大の矩形をiとして繰り返す）

3.のscoreの更新に使用する関数f(score)は2通りあり、IoUの値に対して線形的に更新する方法と指数関数的に更新する方法があります。

1つ目の線形的にscoreを更新する関数f(score)は次式の通りで、矩形iとjのIoUが閾値（IoU_threshold）以上であるなら1-IoUの値を元のスコア値に掛け算した値にスコアが更新されます。

\[
\begin{align}
f(score, iou) =
\begin{cases}
score * (1-iou) \quad (IoU \geq IoU_ threshold) \cr
score \qquad\qquad\qquad (otherwise)
\end{cases}
\end{align}
\]

2つ目の指数関数的にscoreを更新する関数f(score)は次式のようになります。\(\sigma\)はハイパーパラメータで、提案元の論文では0.5等とするようです。調整が必要なパラメータです。

\[
f(score, iou) = score * \exp (-\frac{iou^{2}}{\sigma})
\]

これら2つのスコア更新関数は場合によって使い分けるということを想定しているのかと思います。

上でリンクを貼った提案論文では、物体検出手法R-FCNとFaster-RCNN、データセットMS COCOとPASCAL VOCに対して、スコア更新関数とパラメータIoU_threshold、\(\sigma\)を様々に変更した実験結果が示されており、ケースによって使い分けや調整を行うことが想定されているようです。

また、従来NMSとの比較もこの論文では行われています。

Soft-NMSの実装例

少し長いですが、Soft-NMSのサンプルコードを次に貼り付けます。

# スコアの更新関数（線形）
def f_linear(iou, iou_threshold=0.5):
    if iou >= iou_threshold:
        weight = 1. - iou
    else:
        weight = 1.
    return weight

# スコアの更新関数（指数関数、gauss関数）
def f_gauss(iou, sigma=0.5):
    import math
    return math.exp(-iou*iou/sigma)

def soft_nms(bboxes, scores, classes, \
             iou_threshold=0.5, sigma=0.5, linear=True):
    new_bboxes = [] # Soft-NMS適用後の矩形リスト
    new_scores = [] # Soft-NMS適用後の信頼度(スコア値)リスト
    new_classes = [] # Soft-NMS適用後のクラスのリスト

    while len(bboxes) > 0:
        # スコア最大の矩形のインデックスを取得
        argmax = scores.index(max(scores))

        # スコア最大の矩形、スコア値、クラスをそれぞれのリストから消去
        bbox = bboxes.pop(argmax)
        score = scores.pop(argmax)
        clss = classes.pop(argmax)        

        # スコア最大の矩形と、対応するスコア値、クラスをSoft-NMS適用後のリストに格納
        new_bboxes.append(bbox)
        new_scores.append(score)
        new_classes.append(clss)

        # bboxesに残存する矩形のスコアを更新
        for i, bbox_tmp in enumerate(bboxes):
            # スコア最大の矩形bboxと他の矩形のIoUを計算
            iou_tmp = iou(bbox, bbox_tmp)
            # scoreの値を更新
            if linear: 
                # スコアの更新関数（線形）
                scores[i] = scores[i]*f_linear(iou_tmp, iou_threshold)
            else: 
                # スコアの更新関数（指数関数、gauss関数）
                scores[i] = scores[i]*f_gauss(iou_tmp, sigma)

    return new_bboxes, new_scores, new_classes

関数”f_linear(iou, iou_threshold=0.5)”は、線形的にscoreを更新する関数です。同様に、関数”f_gauss(iou, sigma=0.5)”は、指数関数的にscoreを更新する関数です。

これらは関数”soft_nms(bboxes, scores, classes, iou_threshold=0.5, sigma=0.5, linear=True)”の引数”linear=True”で選択できるようになっています。bboxes, scores, classesは、それぞれ入力となる矩形の座標、スコア、クラスの配列(リスト)で、iou_threshold, sigmaは上で説明したスコア更新関数のハイパーパラメータとなります。

上のサンプルコードの次の部分では、関数”iou(bbox, bbox_tmp)”を使って矩形bboxと矩形bbox_tmpのIoUを計算しています。

# スコア最大の矩形bboxと他の矩形のIoUを計算
iou_tmp = iou(bbox, bbox_tmp)

ここで使用している関数”iou()”はpython標準の関数ではなく、本ブログの過去のエントリで紹介したIoUを計算する関数です。関数”iou()”は次の記事にコードを記載していますので参照ください。

Soft-NMSの適用例(関数soft_nms()の使用例)

上で紹介した関数soft_nms()を使用するには、入力配列(リスト)等の仕様の情報が必要で、その部分が未説明でした。

関数soft_nms()の具体的な使い方について使用例を紹介したいと思うのですが、記事が長くなってしまったので、次の別エントリに記載したいと思います。すみません。

おわりに

ニューラルネットワーク(ディープラーニング)を用いた機械学習による物体検出アルゴリズムの後段処理に使用されるNMSの改良型であるSoft-NMSについて、実装例を紹介しました。

NMSの代わりにSoft-NMSを使用すれば直ちに検出精度が向上するかと言えば、パラメータ調整が必要な面もあり、そうとも言えない部分もあるかと思います。ただ、適切に適用すれば少しmAPなどの検出精度指標が上昇するようです。

どのくらい精度が上がるかは、上で紹介した提soft-NMSの提案論文を参照ください。物体検出AIで、あと少しだけ検出精度を上げたいという場合に、有効な手法であるということは間違いなさそうです。