深度学习_NMS代码详解YOLOv3及Fast R-CNN例子

先放Fast R-CNN的NMS代码

这部分是关于 nms 实现的代码。后续再加下其他的版本。
流程：首先对检测结果的 score 取出最大的元素，然后将置信度最高的框与其他框取交集，计算 iou，将大于 iou 的筛去。然后对剩余的框进行如上操作。

import numpy as np
#这里的NMS是单类别的！多类别则只需要在外加一个for循环遍历每个种类即可
def py_cpu_nms(dets, thresh):
    """Pure Python NMS baseline."""
    x1 = dets[:, 0]
    y1 = dets[:, 1]
    x2 = dets[:, 2]
    y2 = dets[:, 3]
    scores = dets[:, 4]# 取出框的坐标和score

    areas = (x2 - x1 + 1) * (y2 - y1 + 1)# 计算面积
    order = scores.argsort()[::-1] # 用argsort升序，返回的是对应的索引。[::-1]这个呢。
	# a[-1] 取a的最后一个元素  a[:-1] 取除a的最后一个的所有元素
	# a[::-1] 取从后往前的元素。其实就是把scores的索引从大到小排了。
	# a[2::-1] 取a[2]及之前的所有元素，并从a[2]开始返回->a[1]->a[0]

    keep = []
    while order.size > 0:
        i = order[0] # i取最大score的索引

        keep.append(i)
		# 置信度高的预测框即当前框与其他框的交集
		# 选择的区域就是取最大的x1, y1和最小的 x2, y2
        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:]])
		# 计算交叉区域的面积
        w = np.maximum(0.0, xx2 - xx1 + 1)
        h = np.maximum(0.0, yy2 - yy1 + 1)
        inter = w * h
		# 计算IOU,  相交区域 / (当前区域 + 某区域面积 - 相交区域面积)
        ovr = inter / (areas[i] + areas[order[1:]] - inter)
		# 保留IOU小于阈值的框
        inds = np.where(ovr <= thresh)[0]
		# #因为ovr数组的长度比order数组少一个,所以这里要将所有下标后移一位
        order = order[inds + 1]

    return keep

下面是YOLOv3的NMS代码:

def non_max_suppression(prediction, conf_thres=0.4, nms_thres=0.5):
    """
    Removes detections with lower object confidence score than 'conf_thres' and performs
    Non-Maximum Suppression to further filter detections.
    :param prediction: shape为 (batch_size, total_num_bbox, 85)
    :param conf_thres:
    :param nms_thres:
    :return:
    Returns detections with shape:
        (x1, y1, x2, y2, object_conf, class_score, class_pred)
    """
 
    # 把xywh转为x1y1x2y2
    prediction[..., :4] = xywh2xyxy(prediction[..., :4])
 
    # 把output输出化，初始化的结果就是batch_size大小
    output = [None for _ in range(len(prediction))]
 
    # 以每个图片进行处理
    for image_i, image_pred in enumerate(prediction):
        # 把置信度低于阈值的框过滤掉
        image_pred = image_pred[image_pred[:, 4] >= conf_thres]
 
        # 如果过滤后没有剩余的框则直接continue
        if not image_pred.size(0):
            continue
 
        # 计算每个框的置信度：是物体的置信度 乘以 类别预测置信度的最大值
        # prediction的shape为[batch_size x 10647 x 85]
		#10647个预测框，85指 xywh，score，以及80类别
        # img_pred[:, 5:].max(1)[0] 为取出每个框里面类别概率最大的值与置信度相乘
        score = image_pred[:, 4] * image_pred[:, 5:].max(1)[0]
        # 按照score降序
        image_pred = image_pred[(-score).argsort()]
 
        # 得到每一行类别最大的概率与类别代号,keepdim=True:输出的维度和输入维度相同
        # class_confs表示每个框类别置信度的最大值，class_preds表示该最大值对应的序号idx（即类别）
        class_confs, class_preds = image_pred[:, 5:].max(1, keepdim=True)
 
        # cat函数就是按某个维度进行拼接
		#detections:[x1,y1,x2,y2,conf,cls_conf,cls_idx]
		#其中0:3为x1y1x2y2，4为物体置信度score（detections按照score排序）
        # 5为该框中类别置信度的最大值，6为该最大值对应的类别索引
        detections = torch.cat((image_pred[:, :5], class_confs.float(), class_preds.float()), 1)
 
 		# 非极大值抑制
        keep_boxes = []
        while detections.size(0):
            # 看当前列表中的第一个框（即置信度最大的框）和所有剩余框的iou是否大于阀值，如果大于阀值需要考虑是否被丢弃
            # 这里的large_overlap表示重叠overlap过大的索引，这些索引的框可能会被丢弃，这里的阀值取的0.5
			#score最大的框(第一个框)和其他框做交并比，并把大于nms_thres的取出存入large_overlap
            large_overlap = bbox_iou(detections[0, :4].unsqueeze(0), detections[:, :4]) > nms_thres
 
            # 查看是否和第一个框预测的同一个类别标签，如果和第一个框预测的同样类别标签，则考虑是否被丢弃
            label_match = detections[0, -1] == detections[:, -1]
 
            # Indices of boxes with lower confidence scores, large IOUs and matching labels
            # 即同时满足 和第一个框iou大于阀值和类别相同，才丢弃这个框
            invalid = large_overlap & label_match
            weights = detections[invalid, 4:5]  # 要丢弃的框的物体置信度值
            # Merge overlapping bboxes by order of confidence
            # 满足上面条件的框，通过目标置信度对bbox做加权平均处理，用来计算平均框的位置。加权平均合并
            detections[0, :4] = (weights * detections[invalid, :4]).sum(0) / weights.sum()
			# 添加加权平均的detection
            keep_boxes += [detections[0]]
            detections = detections[~invalid] # 取剩余部分的预测值
        if keep_boxes:
		# torch.stack()沿着一个新维度对输入张量序列进行连接。 序列中所有的张量都应该为相同形状。
            output[image_i] = torch.stack(keep_boxes)# 非极大值抑制的结果
 
    return output