nms 实现

序言

nms被用在检测的后处理，目的是剔除多余的检测框。接下来用numpy array实现。

注意点

np.maximum和np.minimum是用来返回对应位置的最大值和最小值。np.argsort是用来返回排序后的位置索引。

代码

import numpy as np
import pdb

def cal_iou(box1, box2):
    xx1 = np.maximum(box1[:, 0], box2[:, 0])
    yy1 = np.maximum(box1[:, 1], box2[:, 1]) 
    xx2 = np.minimum(box1[:, 2], box2[:, 2])
    yy2 = np.minimum(box1[:, 3], box2[:, 3])
    union = (xx2-xx1)*(yy2-yy1)
    area1 = (box1[:, 2]-box1[:, 0])*(box1[:, 3]-box1[:, 1])
    area2 = (box2[:, 2]-box2[:, 0])*(box2[:, 3]-box2[:, 1])
    ious = union/area1+area2-union
    return ious


def nms(detections, nms_thre=0.5):
    """
    detections[[xmin, ymin, xmax, ymax, score],
               [xmin, ymin, xmax, ymax, score],
               [xmin, ymin, xmax, ymax, score],
               ...]
    """
    detections = np.array(detections)
    order = np.argsort(detections[:, -1])[::-1]
    detections = detections[order]
    keep = []
    while len(detections):
        keep.append(detections[0])
        best_box = detections[0][np.newaxis,:]
        ious = cal_iou(best_box, detections)
        valid = ious > nms_thre
        detections = detections[~valid]
        
    keep = np.array(keep)
    return keep


if __name__ == "__main__":
    detections = [[100, 50, 150, 100, 0.7],
                  [120, 60, 150, 100, 0.6],
                  [110, 60, 150, 100, 0.5],
                  [80, 30, 180, 120, 0.9],
                  [110, 60, 150, 100, 0.2],]
    result = nms(detections)
    print(result)

结果

[[ 80.   30.  180.  120.    0.9]
 [100.   50.  150.  100.    0.7]
 [120.   60.  150.  100.    0.6]]