【CV基石】对目标检测原始边框进行NMS

NMS之前

# -*- coding: UTF-8 -*-
import numpy as np
import cv2

def nms(predicts_dict, threshold=0.2):
    """Implement non-maximum supression on predict bounding boxes.
    Parameters
    ----------
        predicts_dict: dict.
            Predict bounding boxes dict:
            {'person': [[x1_1, y1_1, x2_1, y2_1, scores1], [x1_2, y1_2, x2_2, y2_2, scores2], ...], 
             'car': [[x1_1, y1_1, x2_1, y2_1, scores1], [x1_2, y1_2, x2_2, y2_2, scores2], ...]}
        threshhold: float, default is 0.2.
            iou threshold.
    Returns
    -------
        predicts_dict: dict.
            The result of predict bounding boxes processed by non-maximum suppression.
    """
    # 对每一个类别的目标分别进行NMS
    for object_name, bbox in predicts_dict.items():
        bbox_array = np.array(bbox, dtype=np.float) 
        # 获取当前目标类别下所有bounding box的坐标和confidence, 并计算所有bbox的面积
        x1, y1, x2, y2, scores = bbox_array[:,0], bbox_array[:,1], bbox_array[:,2], bbox_array[:,3], bbox_array[:,4]
        areas = (x2 - x1 + 1) * (y2 - y1 + 1)
        print('areas shape = ', areas.shape) 
        # 对当前类别下所有的bbox的confidence进行从高到低排序, order保存索引信息.
        order = scores.argsort()[::-1]
        print('order = ', order)
        # 存放最终保留的bbox的索引信息
        keep = [] 
        # 依次按confidence从高到低遍历bbox，移除所有与该矩形框的IOU值大于threshold的矩形框
        while order.size > 0:
            i = order[0]
            # 保留当前最大confidence对应的bbox索引
            keep.append(i)  
            # 获取所有与当前bbox的交集对应的顶点坐标, 并计算IOU
            # 这里是同时计算当前confidence最大的bbox与剩余其他所有bbox的IOU
            # 当order.size=1时, 下面的计算结果都为np.array([]), 不影响最终结果
            xx1 = np.maximum(x1[i], x1[order[1:]])
            yy1 = np.maximum(y1[i], y1[order[1:]])
            xx2 = np.minimum(x2[i], x2[order[1:]])
            yy2 = np.minimum(y2[i], y2[order[1:]])
            inter = np.maximum(0.0, xx2 - xx1 + 1) * np.maximum(0.0, yy2 - yy1 + 1)
            iou = inter / (areas[i] + areas[order[1:]] - inter)
            print('iou =', iou) 
            # 输出没有被移除的bbox索引, 即相对于iou向量的索引
            print(np.where(iou<=threshold)) 
            # 获取保留下来的索引, 因为没有计算与自身的IOU, 所以索引相差1, 需要加上
            indexs = np.where(iou<=threshold)[0] + 1
            print('indexs = ', type(indexs))
            # 更新保留下来的索引
            order = order[indexs]
            print('order = ', order)
        bbox = bbox_array[keep]
        predicts_dict[object_name] = bbox.tolist()
        # predicts_dict = predicts_dict
    return predicts_dict

def main():
    img_dir = 'image/dog-cycle-car.png'
    img = cv2.imread(img_dir)
    img_nms = img.copy()
    predicts_dict = {'dog': [[83, 178, 240, 416, 0.912366], 
                             [102, 171, 244, 427, 0.981654],
                             [67, 187, 231, 408, 0.281623],
                             [118, 161, 259, 447, 0.671524]], 
                     'bike': [[95, 103, 446, 332, 0.998216], 
                              [109, 76, 428, 344, 0.455442],
                              [87, 96, 456, 328, 0.912163]],
                     'car': [[387, 34, 527, 158, 0.401915],
                             [351, 49, 545, 142, 0.767138],
                             [367, 57, 540, 135, 0.979883]]}
    color_dict = {'dog': (255, 0, 0),
                  'bike': (0, 255, 0),
                  'car': (0, 0, 255)}

    # plot before using NMS
    for object_name, bbox in predicts_dict.items():
        for bb in bbox:
            print(bb[0])
            print(color_dict[object_name])
            cv2.rectangle(img, (int(bb[0]), int(bb[1])), (int(bb[2]), int(bb[3])), color_dict[object_name], 2)
            cv2.putText(img, '%s:%.3f' % (object_name, bb[4]), (int(bb[0]), int(bb[1])), cv2.FONT_HERSHEY_SIMPLEX, 1, (0,255,255), 2)

    rst_dict = nms(predicts_dict)
    
    # plot after using NMS
    for object_name, bbox in rst_dict.items():
        for bb in bbox:
            cv2.rectangle(img_nms, (int(bb[0]), int(bb[1])), (int(bb[2]), int(bb[3])), color_dict[object_name], 2)
            cv2.putText(img_nms, '%s:%.3f' % (object_name, bb[4]), (int(bb[0]), int(bb[1])), cv2.FONT_HERSHEY_SIMPLEX, 1, (0,255,255), 2)

    cv2.imwrite('image/nms_before.png', img)
    cv2.imwrite('image/nms_after.png', img_nms)
    cv2.imshow('nms_before',img)
    cv2.imshow('nms_after',img_nms)
    cv2.waitKey(0)

if __name__ == "__main__":
    main()

NMS之后

改进

Soft-NMS

---

Technical Exchange