1、前言

目前yolov5使用的是NMS进行极大值抑制,本篇文章是要将各类NMS添加到yolov5中,同时可以使用不同的IOU进行预测框处理。

NMS概念
NMS(Non-maximum suppression)是非极大值抑制, 目的是过滤掉重复的框。 为了保证检测的准确性, 检测网络的输出框一般都比较密集, 对一个物体, 会有多个预测框,NMS就是为了过滤掉这些重复的框, 保留质量最好的那一个框。如下图,设置的阈值可能为0.6,将概率大于0.6的预测框都保留下来,再进行nms非极大值抑制处理,保留与标注框最接近的预测框。
在这里插入图片描述

Soft-NMS概念
Soft-NMS认为在进行NMS的时候要同时考虑重合程度和得分。如果存在同一类别的两个目标彼此重叠,即使两个预测框的得分均较高,在NMS下因两个预测框的重合程度较大也只能保留一个,如下图:

在这里插入图片描述

两个预测框均是针对马的,两个预测分数(0.95、0.80)也均较高,重叠程度(IoU)很可能大于阈值,如果执行NMS,仅可能保留红框,Soft-NMS为避免该问题,提出根据与最大得分框的IoU降低与之重叠框的得分而非直接去除的方法。

这样,如果两个高得分的框重叠较多,在得分最高的框被保留后,另一个分数同样较高的框虽然得分有所降低,但因其本身得分相比其它框更高,在下一轮比较中有希望脱颖而出。

2、各类NMS代码实现

2.1、general.py

(1)在utils\general.py文件,找到non_max_suppression函数

在这里插入图片描述

(2)在该函数上方添加下方代码

def box_iou_for_nms(box1, box2, GIoU=False, DIoU=False, CIoU=False, SIoU=False, EIou=False, eps=1e-7):
    # Returns Intersection over Union (IoU) of box1(1,4) to box2(n,4)

    b1_x1, b1_y1, b1_x2, b1_y2 = box1.chunk(4, -1)
    b2_x1, b2_y1, b2_x2, b2_y2 = box2.chunk(4, -1)
    w1, h1 = b1_x2 - b1_x1, (b1_y2 - b1_y1).clamp(eps)
    w2, h2 = b2_x2 - b2_x1, (b2_y2 - b2_y1).clamp(eps)

    # Intersection area
    inter = (b1_x2.minimum(b2_x2) - b1_x1.maximum(b2_x1)).clamp(0) * \
            (b1_y2.minimum(b2_y2) - b1_y1.maximum(b2_y1)).clamp(0)

    # Union Area
    union = w1 * h1 + w2 * h2 - inter + eps

    # IoU
    iou = inter / union
    cw = b1_x2.maximum(b2_x2) - b1_x1.minimum(b2_x1)  # convex (smallest enclosing box) width
    ch = b1_y2.maximum(b2_y2) - b1_y1.minimum(b2_y1)  # convex height
    if CIoU or DIoU or GIoU or EIou:
        if CIoU or DIoU or EIou:  # Distance or Complete IoU https://arxiv.org/abs/1911.08287v1
            c2 = cw ** 2 + ch ** 2 + eps  # convex diagonal squared
            rho2 = ((b2_x1 + b2_x2 - b1_x1 - b1_x2) ** 2 + (b2_y1 + b2_y2 - b1_y1 - b1_y2) ** 2) / 4  # center dist ** 2
            if CIoU:  # https://github.com/Zzh-tju/DIoU-SSD-pytorch/blob/master/utils/box/box_utils.py#L47
                v = (4 / math.pi ** 2) * (torch.atan(w2 / h2) - torch.atan(w1 / h1)).pow(2)
                with torch.no_grad():
                    alpha = v / (v - iou + (1 + eps))
                return iou - (rho2 / c2 + v * alpha)  # CIoU
            elif EIou:
                rho_w2 = ((b2_x2 - b2_x1) - (b1_x2 - b1_x1)) ** 2
                rho_h2 = ((b2_y2 - b2_y1) - (b1_y2 - b1_y1)) ** 2
                cw2 = cw ** 2 + eps
                ch2 = ch ** 2 + eps
                return iou - (rho2 / c2 + rho_w2 / cw2 + rho_h2 / ch2)
            return iou - rho2 / c2  # DIoU
        c_area = cw * ch + eps  # convex area
        return iou - (c_area - union) / c_area  # GIoU https://arxiv.org/pdf/1902.09630.pdf
    elif SIoU:
        # SIoU Loss https://arxiv.org/pdf/2205.12740.pdf
        s_cw = (b2_x1 + b2_x2 - b1_x1 - b1_x2) * 0.5 + eps
        s_ch = (b2_y1 + b2_y2 - b1_y1 - b1_y2) * 0.5 + eps
        sigma = torch.pow(s_cw ** 2 + s_ch ** 2, 0.5)
        sin_alpha_1 = torch.abs(s_cw) / sigma
        sin_alpha_2 = torch.abs(s_ch) / sigma
        threshold = pow(2, 0.5) / 2
        sin_alpha = torch.where(sin_alpha_1 > threshold, sin_alpha_2, sin_alpha_1)
        angle_cost = torch.cos(torch.arcsin(sin_alpha) * 2 - math.pi / 2)
        rho_x = (s_cw / cw) ** 2
        rho_y = (s_ch / ch) ** 2
        gamma = angle_cost - 2
        distance_cost = 2 - torch.exp(gamma * rho_x) - torch.exp(gamma * rho_y)
        omiga_w = torch.abs(w1 - w2) / torch.max(w1, w2)
        omiga_h = torch.abs(h1 - h2) / torch.max(h1, h2)
        shape_cost = torch.pow(1 - torch.exp(-1 * omiga_w), 4) + torch.pow(1 - torch.exp(-1 * omiga_h), 4)
        return iou - 0.5 * (distance_cost + shape_cost)
    return iou  # IoU



def soft_nms(bboxes, scores, iou_thresh=0.5, sigma=0.5, score_threshold=0.25):
    order = torch.arange(0, scores.size(0)).to(bboxes.device)
    keep = []

    while order.numel() > 1:
        if order.numel() == 1:
            keep.append(order[0])
            break
        else:
            i = order[0]
            keep.append(i)

        # 修改成你想要使用的IOU,如果不修改默认的就是普通iou
        iou = box_iou_for_nms(bboxes[i], bboxes[order[1:]], CIoU=True).squeeze()

        idx = (iou > iou_thresh).nonzero().squeeze()
        if idx.numel() > 0:
            iou = iou[idx]
            newScores = torch.exp(-torch.pow(iou, 2) / sigma)
            scores[order[idx + 1]] *= newScores

        newOrder = (scores[order[1:]] > score_threshold).nonzero().squeeze()
        if newOrder.numel() == 0:
            break
        else:
            maxScoreIndex = torch.argmax(scores[order[newOrder + 1]])
            if maxScoreIndex != 0:
                newOrder[[0, maxScoreIndex],] = newOrder[[maxScoreIndex, 0],]
            order = order[newOrder + 1]

    return torch.LongTensor(keep)

(3)non_max_suppression函数里的原始NMS,将其注释掉,添加如下代码:

i = soft_nms(boxes, scores, iou_thres)

在这里插入图片描述

3、各类NMS实现

utils\general.py找到刚刚添加的soft_nms函数里面找到下面这行代码,将你想要使用的Iou修改为True,如果不修改默认的就是普通iou,添加IOU之后就能训练模型了。

3.1、Soft-NMS

在这里插入图片描述

3.2、GIoU-NMS

iou = box_iou_for_nms(bboxes[i], bboxes[order[1:]], GIoU=True).squeeze()

3.3、DIoU-NMS

iou = box_iou_for_nms(bboxes[i], bboxes[order[1:]], DIoU=True).squeeze()

3.4、CIoU-NMS

iou = box_iou_for_nms(bboxes[i], bboxes[order[1:]], CIoU=True).squeeze()

3.5、EIoU-NMS

iou = box_iou_for_nms(bboxes[i], bboxes[order[1:]], EIoU=True).squeeze()

4、目标检测系列文章

  1. YOLOv5s网络模型讲解(一看就会)
  2. 生活垃圾数据集(YOLO版)
  3. YOLOv5如何训练自己的数据集
  4. 双向控制舵机(树莓派版)
  5. 树莓派部署YOLOv5目标检测(详细篇)
  6. YOLO_Tracking 实践 (环境搭建 & 案例测试)
  7. 目标检测:数据集划分 & XML数据集转YOLO标签
  8. DeepSort行人车辆识别系统(实现目标检测+跟踪+统计)
  9. YOLOv5参数大全(parse_opt篇)
  10. YOLOv5改进(一)-- 轻量化YOLOv5s模型
  11. YOLOv5改进(二)-- 目标检测优化点(添加小目标头检测)
  12. YOLOv5改进(三)-- 引进Focaler-IoU损失函数
  13. YOLOv5改进(四)–轻量化模型ShuffleNetv2
  14. YOLOv5改进(五)-- 轻量化模型MobileNetv3
  15. YOLOv5改进(六)–引入YOLOv8中C2F模块
  16. YOLOv5改进(七)–改进损失函数EIoU、Alpha-IoU、SIoU、Focal-EIOU
Logo

开放原子开发者工作坊旨在鼓励更多人参与开源活动,与志同道合的开发者们相互交流开发经验、分享开发心得、获取前沿技术趋势。工作坊有多种形式的开发者活动,如meetup、训练营等,主打技术交流,干货满满,真诚地邀请各位开发者共同参与!

更多推荐