PythonのnumpyでNMSを高速に演算する(20倍高速?)

重複した矩形を統合するアルゴリズムNMS (Non-Maximum Suppression)を、numpyで高速に計算する方法を紹介します。numpyを使わずpythonリスト(list)を使用する実装と比べて約20倍に高速化できました。ただし、高速化の度合いはPC環境等によって変わります。

この記事でできること

SSDやYOLOといった物体検出AIでは、人の顔などの検出目標の周辺にたくさんの矩形が出力されてしまいます。それらの矩形をマージしてすっきりさせるためのアルゴリズムがNMS (Non-Maximum Suppression)です。

この記事では、numpyのブロードキャスト演算を使って高速にNMSを計算する方法を紹介します。

一見同様のサンプルコードが「(Faster) Non-Maximum Suppression in Python(英語記事)」で紹介されていますが、この記事で計算しているIoU (Intersection over Union)は一般的に使用されるものとは異なるようです。Intersectionに対して2つの矩形(長方形)のUnionの面積で割るのではなく、片方の矩形(長方形)の面積で割るというソースコードになっています。

この記事では、一般的に使用されるIoUのアルゴリズムで高速にNMSを計算する実装例を紹介します。また、計算時間の計測結果も記載します。

以前紹介した普通のpythonリストを用いたNMSの実装例に比べ、約20倍高速化できました (私の環境での値になります)。

処理時間の計測では、下の画像のようにランダムに矩形を生成して(1か所あたり100個)、各回で乱数を変えて生成した矩形に対して100回NMSを実行した処理時間を測りました。

NMS適用前のランダムに付与した矩形画像
NMS適用前のランダムに付与した矩形画像

上の画像に対してNMSを適用した結果が次のようになります。少し削除しきれていない矩形はありますが、NMSは万能ではなく仕方ないところです。

高速なNMS適用後の矩形画像
高速なNMS適用後の矩形画像

NMS(Non-Maximum Suppression)について

この記事では、numpyのブロードキャスト演算を使って複数の矩形のIoUを一気に計算したり、IoUの値の閾値判定を一度に行うことで高速にNMSを実行する方法を紹介します。

NMSとは何かということやNMSの適用例は、次の記事で紹介しました。この過去記事ではnumpyを使わず、リスト(list)を使ってfor文で1つずつ計算し、NMSを行う実装例を紹介しています。

numpyによる高速なNMSの実装例

高速にNMSを計算するサンプルコードは次の通りです。関数”iou_np()”で1つの矩形と複数の矩形のIoUを一気に計算しています。関数”nms_fast()”がNMSを計算するコードの本体です。

# 矩形aと、複数の矩形bのIoUを計算
def iou_np(a, b, a_area, b_area):
    # aは1つの矩形を表すshape=(4,)のnumpy配列
    # array([xmin, ymin, xmax, ymax])
    # bは任意のN個の矩形を表すshape=(N, 4)のnumpy配列
    # 2次元目の4は、array([xmin, ymin, xmax, ymax])
    
    # a_areaは矩形aの面積
    # b_areaはbに含まれる矩形のそれぞれの面積
    # shape=(N,)のnumpy配列。Nは矩形の数
    
    # aとbの矩形の共通部分(intersection)の面積を計算するために、
    # N個のbについて、aとの共通部分のxmin, ymin, xmax, ymaxを一気に計算
    abx_mn = np.maximum(a[0], b[:,0]) # xmin
    aby_mn = np.maximum(a[1], b[:,1]) # ymin
    abx_mx = np.minimum(a[2], b[:,2]) # xmax
    aby_mx = np.minimum(a[3], b[:,3]) # ymax
    # 共通部分の幅を計算。共通部分が無ければ0
    w = np.maximum(0, abx_mx - abx_mn + 1)
    # 共通部分の高さを計算。共通部分が無ければ0
    h = np.maximum(0, aby_mx - aby_mn + 1)
    # 共通部分の面積を計算。共通部分が無ければ0
    intersect = w*h
    
    # N個のbについて、aとのIoUを一気に計算
    iou_np = intersect / (a_area + b_area - intersect)
    return iou_np

# NMSの計算
def nms_fast(bboxes, scores, classes, iou_threshold=0.5):
    # bboxesは任意のN個の矩形を格納したshape=(N, 4)のnumpy配列
    # 2次元目の4要素は、array([xmin, ymin, xmax, ymax])
    # scoresは任意のN個の信頼度を格納したshape=(N,)のnumpy配列
    # classesは任意のN個のクラスを格納したshape=(N,)のnumpy配列
    
    # bboxesの矩形の面積を一気に計算
    areas = (bboxes[:,2] - bboxes[:,0] + 1) \
             * (bboxes[:,3] - bboxes[:,1] + 1)
    
    # scoreの昇順(小さい順)の矩形インデックスのリストを取得
    sort_index = np.argsort(scores)
    
    i = -1 # 未処理の矩形のindex
    while(len(sort_index) >= 2 - i):
        # score最大のindexを取得
        max_scr_ind = sort_index[i]
        # score最大以外のindexを取得
        ind_list = sort_index[:i]
        # score最大の矩形それ以外の矩形のIoUを計算
        iou = iou_np(bboxes[max_scr_ind], bboxes[ind_list], \
                     areas[max_scr_ind], areas[ind_list])
        
        # IoUが閾値iou_threshold以上の矩形を計算
        del_index = np.where(iou >= iou_threshold)
        # IoUが閾値iou_threshold以上の矩形を削除
        sort_index = np.delete(sort_index, del_index)
        #print(len(sort_index), i, flush=True)
        i -= 1 # 未処理の矩形のindexを1減らす
    
    # bboxes, scores, classesから削除されなかった矩形のindexのみを抽出
    bboxes = bboxes[sort_index]
    scores = scores[sort_index]
    classes = classes[sort_index]
    
    return bboxes, scores, classes

NMSの計算に必要な、1つの矩形に対する複数の矩形のIoUそれぞれを一気に計算する方法や実装例は次の過去記事でも紹介しました。この過去記事で紹介している関数”iou_np()”と今回の関数”iou_np()”は少し異なり、矩形の面積”a_area”、”b_area”を引数として使用しています。これらは何度も使用するので、最初に一回計算した値を引数としているわけです。

その他、細かい処理の内容はサンプルコード内のコメント文を参考にしてみてください。

numpyによる高速なNMSの適用例

関数”nms_fast()”を実行するサンプルコード

上に記載した関数”nms_fast()”の適用例を紹介します。今回使用するサンプルコードは次の通りです。このコード内で定義している関数”random_rects()”では、呼び出されるとランダムに矩形を100個生成して返します。

def random_rects(bboxes, scores, classes):
    rnd_bboxes = []
    rnd_scores = []
    rnd_classes = []
    for i, bbox in enumerate(bboxes):
        for dmy in range(100): # 100個の矩形を乱数で生成
            # 上で作成した矩形座標rectsを-15~+15範囲でずらした矩形を乱数で生成
            rnd_bbox = [a - 15 + 30*random.random() for a in bbox]
            rnd_bboxes.append(rnd_bbox)
            
            # scoreを0~0.9の間の一様乱数で生成
            scr = random.random()*0.9
            rnd_scores.append(scr)
            
            # 確率0.5でman/womanのクラス生成
            if random.random() < 0.5: cl = "man"
            else: cl = "woman"
            rnd_classes.append(cl)
    
    rnd_bboxes.extend(bboxes)
    rnd_classes.extend(classes)
    rnd_scores.extend(scores)
    return rnd_bboxes, rnd_scores, rnd_classes

if __name__ == '__main__':
    import random
    import time
    
    iou_threshold = 0.3
    # 顔の矩形リスト
    bboxes = [[161,50,260,161], 
              [318,135,418,243], 
              [46,267,142,373], 
              [181,244,275,352], 
              [340,351,440,455], 
              [500,257,598,371], 
              [648,33,742,146]]
    classes = ["woman", "woman", "man", "man", "woman", "man", "man"]
    # 信頼度の値は適当
    scores = [0.91, 0.92, 0.93, 0.94, 0.95, 0.96, 0.97]

    random.seed(314) # 乱数シードを設定
    time_list = [] # 処理時間を格納するリスト(list)
    for dmy in range(100): # 100回実行
        # NMSを計算するための矩形をランダムに生成
        rnd_bboxes, rnd_scores, rnd_classes\
            = random_rects(bboxes, scores, classes)
        bboxes_np = np.array(rnd_bboxes)
        scores_np = np.array(rnd_scores)
        classes_np = np.array(rnd_classes)
        start = time.perf_counter() # 時間計測の関数
        # NMSを実行
        nms_bboxes, nms_scores, nms_classes = \
            nms_fast(bboxes_np, scores_np, classes_np, \
                     iou_threshold=iou_threshold)
        nms_time = time.perf_counter() - start # 時間計測の関数
        # 計測した時間(秒)を1000倍して単位をミリ秒をして、time_listに追加
        time_list.append(nms_time*1000)
    
    # 処理時間の平均を取得
    print ("NMS time: {0}".format(sum(time_list)/len(time_list)) + " [msec]")

このサンプルコードでは、まず複数の矩形とそれに対応する信頼度とクラスをスクリプトに直書きし、リストbboxes, scores, classesに格納します。

その後、関数”random_rects()”を使ってbboxes内の各矩形を基準にランダムに100個ずつ矩形を生成した後、関数”nms_fast()”でNMSを実行するということを、for文で100回繰り返しています。

for文内の各繰り返しで、関数”nms_fast()”の処理時間を測定しています。測定はサンプルコードの次の部分で行っています。


start = time.perf_counter() # 時間計測の関数

## NMSの処理

nms_time = time.perf_counter() - start # 時間計測の関数
# 計測した時間(秒)を1000倍して単位をミリ秒をして、time_listに追加
time_list.append(nms_time*1000)

詳しくは次の記事で説明していますので参照ください。

また、次の1文は乱数シードを変更する処理です。

random.seed(314) # 乱数シードを設定

疑似乱数を使用する際、乱数シードの設定は重要となります。設定方法の詳細は次の過去記事で説明しています。

関数”nms_fast()”を適用した画像例

上のサンプルコードの実行例を画像で示します。

次の画像は、関数”random_rects()”で矩形をランダムに100個ずつ顔に付与した画像です。NMS適用前になります。

NMS適用前のランダムに付与した矩形画像
NMS適用前のランダムに付与した矩形画像

矩形を100個も重畳すると大変なことになりますね。これに対してNMSを適用した結果を画像化すると次のようになります。

高速なNMS適用後の矩形画像
高速なNMS適用後の矩形画像

これらは、この記事の上の方で紹介した画像と同じです。

画像に対して、矩形(長方形)を描画したり、”man”や”0.93″といったクラス名や信頼度数値のテキストを描画したりする方法は以前紹介しましたので参考にしてみてください。

計算時間計測結果

上で紹介したサンプルコードでは、100回のNMS計算の平均処理時間をprint文で出力するようにしています。私の環境では次のようになりました。

NMS time: 31.734449999998592 [msec]

1回あたり約32ミリ秒(msec)でした。

過去の記事「pythonでNMSを実装し、複数の矩形をマージする」で紹介した、リスト(list)を使ってfor文で1つずつIoUを計算するサンプルコードを使用して同じようにNMSを実行し、処理時間を計測した結果は次のようになりました。

NMS time: 620.6073499999998 [msec]

1回あたり約620ミリ秒(msec)となり、今回紹介したnumpyのブロードキャスト演算を使用した高速なNMS演算は約20倍高速化できたことがわかります。

ただし、高速化の度合いはPC環境やNMSに適用する矩形(長方形)によって変わりますのでご注意ください。

また、100回実行した処理時間をヒストグラム化したものが次の画像です。横軸が処理時間[ミリ秒(msec)]で、橙色の棒グラフが今回記事、青色の棒グラフが以前紹介した方法の処理時間となります。

100回NMSを実行したときの計算時間のヒストグラム (横軸ミリ秒 [msec]) 橙色が高速なNMS、青色が従来のNMS
100回NMSを実行したときの計算時間のヒストグラム (横軸ミリ秒 [msec])
橙色が高速なNMS、青色が従来のNMS

従来のNMSは600ミリ秒(msec)近辺、高速なNMSは約30ミリ秒(msec)周辺に分布していることがわかります。この記事で紹介した高速なNMSのサンプルコードは、あまりばらつきなく安定的に高速化されていることがわかります。

pythonのmatplotlibライブラリを用いたヒストグラムの作成方法は次の記事で紹介しています。

おわりに

SSDやYOLOといった物体検出AIの後処理として使用される、重複した矩形をマージするアルゴリズムNMSの高速計算方法を紹介しました。numpyのブロードキャスト演算を使って高速化し、numpyを使わずpythonリスト(list)を使用する実装と比べて約20倍に高速化できました。(高速化の度合いは環境などによって変わります)

NMSの改良型で検出性能を向上できるSoft-NMSも同様に高速化しましたので、こちらも参照ください。

コメント

タイトルとURLをコピーしました