一种非极大值抑制（non_max_suppression, nms）的代码实现方式

1. 简介

• 非极大值抑制，non_max_suppression，简称nms，常用于目标检测的后处理，去除多余的检测框。
• 流程大致是：根据某个类别，按照检测框的置信度从大到小排序，选择置信度最高的检测框记为A，计算A与剩余的检测框（B1、B2、… BN）的iou值，若iou值大于设置的阈值，则将B1、B2、… BN中对应的检测框去掉，如此重复操作。

2. 代码

• 结合yolov5模型输出，实现nms的代码。
• 在yolov5中，模型的输出记为pred，（这里pred是模型的直接输出，还没有经过后处理），pred的shape为：(batch_size, num_bbox, 4 + 1 + num_classes)，其中batch_size表示输入的图片数量，num_bbox表示检测到的矩形框数量，4表示4个坐标，1表示检测框的置信度，num_classes表示每个类别的分数。
• pred的检测框坐标格式为：xywh，即检测框的中心坐标以及它的宽高

2.1 坐标形式转换

# coding=utf-8
import numpy as np

def xywh2xyxy(x):
    y = np.copy(x)
    y[..., 0] = x[..., 0] - x[..., 2] / 2  # top left x
    y[..., 1] = x[..., 1] - x[..., 3] / 2  # top left y
    y[..., 2] = x[..., 0] + x[..., 2] / 2  # bottom right x
    y[..., 3] = x[..., 1] + x[..., 3] / 2  # bottom right y
    return y

2.2 iou计算

def cal_iou(det1, det2):
    det1_x1, det1_y1 = det1[..., 0], det1[..., 1]
    det1_x2, det1_y2 = det1[..., 2], det1[..., 3]

    det2_x1, det2_y1 = det2[..., 0], det2[..., 1]
    det2_x2, det2_y2 = det2[..., 2], det2[..., 3]

    x1 = np.maximum(det1_x1, det2_x1)
    y1 = np.maximum(det1_y1, det2_y1)
    x2 = np.minimum(det1_x2, det2_x2)
    y2 = np.minimum(det1_y2, det2_y2)

    area_det1 = (det1_y2 - det1_y1 + 1) * (det1_x2 - det1_x1 + 1)
    area_det2 = (det2_y2 - det2_y1 + 1) * (det2_x2 - det2_x1 + 1)
    inter = np.maximum(0, (y2 - y1 + 1)) * np.maximum(0, (x2 - x1 + 1))

    ious = inter / (area_det1 + area_det2 - inter)
    return ious

2.3 nms

def nms(detections, conf_thres=0.4, nms_thres=0.5):
    outputs = []
    detections = xywh2xyxy(detections)               # 检测框格式转换 xywh -> xyxy
    detections[..., 5:] *= detections[..., 4:5]      # 类别的分数是用obj_conf * cls_conf
    num_classes = detections.shape[2] - 5            # detections shape: (batch_size, num_bbox, 4 + 1 + num_classes)
    candidates = detections[..., 4] > conf_thres     # 用于后续的过滤筛选
    for img_idx, dets in enumerate(detections):
        output = []
        dets = dets[candidates[img_idx]]             # 根据conf_thres过滤
        indexes = dets[..., 4].argsort()[::-1]       # 根据obj_conf大小排序
        dets = dets[indexes]
        classes_dets = dets[..., 5:].argmax(axis=1)  # 每个检测框的类别
        for cls in range(num_classes):               # 按类别循环遍历
            dets_cls = dets[classes_dets == cls]     # shape: (num_bbox, 4 + 1 + num_classes)
            while len(dets_cls):
                det_select = dets_cls[0]
                # det_select shape: (x1, y1, x2, y2, score, label)
                det_select = np.concatenate((det_select[..., :4], [det_select[..., 5:].max(), cls]))
                output.append(det_select)
                dets_cls = dets_cls[1:]
                if len(dets_cls):
                    ious = cal_iou(det_select, dets_cls)
                    indexes = np.where(ious <= nms_thres)[0]
                    dets_cls = dets_cls[indexes]
        outputs.append(output)
    return np.array(outputs)

这里得到的outputs还需要映射回原图像，以下图为例：

• 如图所示，原始图像经过resize之后得到红色框，红色框经过padding得到黑色框，绿色框表示检测框在黑色框中的位置（即以黑色框左上角为原点的位置）。
• 将检测结果映射回原始图像时，要先去除padding的影响，再去除resize的影响，这样就得到检测框在原图中的位置。具体步骤如下：

1. 检测框先向上移动t个像素，再向左移动p个像素，其中t、p表示在上方向、左方向padding的像素个数，具体操作将x1 = x1 - p, x2 = x2 - p, y1 = y1 - t, y2 = y2 - t; 这样就得到检测框在没有padding操作下的位置（即以红色框左上角为原点的位置）；
2. 去除resize的影响。只要将坐标除以图像resize后与原图像的的比例gain即可，假设原始图像高、宽为ori_h, ori_w，resize之后的宽高为new_h, new_w，则ratio = min(new_h/ori_h, new_w/ori_w), 然后让x1 = x1/ratio, y1 = y1/ratio, x2 = x2/ratio, y2 = y2/ratio，即得到原始图像的位置。
3. 代码如下:

def scale_coords(output, ori_shape, ratio=1.0, top_pad=0, left_pad=0):
    ori_h, ori_w = ori_shape
    output[:, [0, 2]] -= left_pad
    output[:, [1, 3]] -= top_pad
    output[:, :4] /= ratio
    output[:, [0, 2]] = output[:, [0, 2]].clip(0, ori_w)  # x1, x2
    output[:, [1, 3]] = output[:, [1, 3]].clip(0, ori_h)  # y1, y2

完成。