OpenCV – 物体検出で使われる Non Maximum Suppression について

import cv2
import numpy as np
from IPython.display import Image, display
from matplotlib import pyplot as plt


def imshow(img):
    """ndarray 配列をインラインで Notebook 上に表示する。
    """
    ret, encoded = cv2.imencode(".jpg", img)
    display(Image(encoded))


# 入力画像、テンプレート画像を読み込む。
img = cv2.imread("sample.jpg")  # 入力画像
templ = cv2.imread("template.jpg")  # テンプレート画像

# テンプレートマッチングを行う。
result = cv2.matchTemplate(img, templ, cv2.TM_CCOEFF_NORMED)

# 類似度が 0.6 以上の位置及びスコアを取得する。
positions = np.where(result >= 0.9)
scores = result[positions]

# 各要素が (x1, y1, x2, y2) である短形一覧を作成する。
boxes = []
h, w = templ.shape[:2]  # テンプレート画像の高さ及び幅
for y, x in zip(*positions):
    boxes.append([x, y, x + w - 1, y + h - 1])
boxes = np.array(boxes)

def draw_boxes(img, boxes):
    dst = img.copy()
    for x1, y1, x2, y2 in boxes:
        cv2.rectangle(
            dst, (x1, y1), (x2, y2), color=(0, 255, 0), thickness=2,
        )
    imshow(dst)
    print("number of boxes", len(boxes))


# 描画する。
draw_boxes(img, boxes)

def non_max_suppression(boxes, scores, overlap_thresh):
    """Non Maximum Suppression (NMS) を行う。

    Args:
        boxes: (N, 4) の numpy 配列。矩形の一覧。
        overlap_thresh: [0, 1] の実数。閾値。

    Returns:
        boxes : (M, 4) の numpy 配列。Non Maximum Suppression により残った矩形の一覧。
    """
    if len(boxes) <= 1:
        return boxes

    # float 型に変換する。
    boxes = boxes.astype("float")

    # (NumBoxes, 4) の numpy 配列を x1, y1, x2, y2 の一覧を表す4つの (NumBoxes, 1) の numpy 配列に分割する。
    x1, y1, x2, y2 = np.squeeze(np.split(boxes, 4, axis=1))

    # 矩形の面積を計算する。
    area = (x2 - x1 + 1) * (y2 - y1 + 1)

    indices = np.argsort(scores)  # スコアを降順にソートしたインデックス一覧
    selected = []  # NMS により選択されたインデックス一覧

    # indices がなくなるまでループする。
    while len(indices) > 0:
        # indices は降順にソートされているので、一番最後の要素の値 (インデックス) が
        # 残っている中で最もスコアが高い。
        last = len(indices) - 1

        selected_index = indices[last]
        remaining_indices = indices[:last]
        selected.append(selected_index)

        # 選択した短形と残りの短形の共通部分の x1, y1, x2, y2 を計算する。
        i_x1 = np.maximum(x1[selected_index], x1[remaining_indices])
        i_y1 = np.maximum(y1[selected_index], y1[remaining_indices])
        i_x2 = np.minimum(x2[selected_index], x2[remaining_indices])
        i_y2 = np.minimum(y2[selected_index], y2[remaining_indices])

        # 選択した短形と残りの短形の共通部分の幅及び高さを計算する。
        # 共通部分がない場合は、幅や高さは負の値になるので、その場合、幅や高さは 0 とする。
        i_w = np.maximum(0, i_x2 - i_x1 + 1)
        i_h = np.maximum(0, i_y2 - i_y1 + 1)

        # 選択した短形と残りの短形の Overlap Ratio を計算する。
        overlap = (i_w * i_h) / area[remaining_indices]

        # 選択した短形及び OVerlap Ratio が閾値以上の短形を indices から削除する。
        indices = np.delete(
            indices, np.concatenate(([last], np.where(overlap > overlap_thresh)[0]))
        )

    # 選択された短形の一覧を返す。
    return boxes[selected].astype("int")

# Non Maximum Suppression を行う。
boxes = non_max_suppression(boxes, scores, overlap_thresh=0.6)

# NMS 後に残った短形一覧を描画する。
draw_boxes(img, boxes)