目录
- 1. 简介
- 2. 代码
* - 2.1 坐标形式转换
- 2.2 iou计算
-
简介
-
非极大值抑制，non_max_suppression，简称nms，常用于目标检测的后处理，去除多余的检测框。
-
流程大致是：根据某个类别，按照检测框的置信度从大到小排序，选择置信度最高的检测框记为A，计算A与剩余的检测框（B1、B2、… BN）的iou值，若iou值大于设置的阈值，则将B1、B2、… BN中对应的检测框去掉，如此重复操作。
-
代码
-
结合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 坐标形式转换
import numpy as np
def xywh2xyxy(x):
y = np.copy(x)
y[..., 0] = x[..., 0] - x[..., 2] / 2
y[..., 1] = x[..., 1] - x[..., 3] / 2
y[..., 2] = x[..., 0] + x[..., 2] / 2
y[..., 3] = x[..., 1] + x[..., 3] / 2
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)
detections[..., 5:] *= detections[..., 4:5]
num_classes = detections.shape[2] - 5
candidates = detections[..., 4] > conf_thres
for img_idx, dets in enumerate(detections):
output = []
dets = dets[candidates[img_idx]]
indexes = dets[..., 4].argsort()[::-1]
dets = dets[indexes]
classes_dets = dets[..., 5:].argmax(axis=1)
for cls in range(num_classes):
dets_cls = dets[classes_dets == cls]
while len(dets_cls):
det_select = dets_cls[0]
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的影响，这样就得到检测框在原图中的位置。具体步骤如下：
-
检测框先向上移动t个像素，再向左移动p个像素，其中t、p表示在上方向、左方向padding的像素个数，具体操作将x1 = x1 – p, x2 = x2 – p, y1 = y1 – t, y2 = y2 – t; 这样就得到检测框在没有padding操作下的位置（即以红色框左上角为原点的位置）；
- 去除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，即得到原始图像的位置。
- 代码如下:
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)
output[:, [1, 3]] = output[:, [1, 3]].clip(0, ori_h)
完成。
Original: https://blog.csdn.net/weixin_43508499/article/details/125884082
Author: lokvke
Title: 一种非极大值抑制（non_max_suppression, nms）的代码实现方式
原创文章受到原创版权保护。转载请注明出处：https://www.johngo689.com/762934/
转载文章受原作者版权保护。转载请注明原作者出处！