👀🎉📜系列文章目录
【YOLOv5改进系列(1)】高效涨点----使用EIoU、Alpha-IoU、SIoU、Focal-EIOU替换CIou
【YOLOv5改进系列(2)】高效涨点----Wise-IoU详细解读及使用Wise-IoU(WIOU)替换CIOU
【YOLOv5改进系列(3)】高效涨点----Optimal Transport Assignment:OTA最优传输方法
【YOLOv5改进系列(4)】高效涨点----添加可变形卷积DCNv2
【YOLOv5改进系列(5)】高效涨点----添加密集小目标检测NWD方法
【YOLOv5改进系列(6)】高效涨点----使用DAMO-YOLO中的Efficient RepGFPN模块替换yolov5中的Neck部分
【YOLOv5改进系列(7)】高效涨点----使用yolov8中的C2F模块替换yolov5中的C3模块
【YOLOv5改进系列(8)】高效涨点----添加yolov7中Aux head 辅助训练头
【YOLOv5改进系列(9)】高效涨点----使用CAM(上下文增强模块)替换掉yolov5中的SPPF模块
【YOLOv5改进系列(10)】高效涨点----将CAM(上下文增强模块)添加到Neck特征融合模块当中
🚀🚀🚀前言
🔥🔥🔥目前yolov5使用的是NMS进行极大值抑制,本篇文章是要将soft-nms添加到yolov5中,同时可以使用不同的IOU进行预测框处理,在看本文之前强烈推荐看一下改进系列(1)和(2)这两个讲了EIoU、Alpha-IoU、SIoU、Focal-EIOU、Wise-IoU
的相关概念,以及如何使用不同的Iou替换掉yolov5中的CIOU。我自己也在钢轨表面疵点数据集进行测试,发现效果还是有一些,但是训练速度比较慢,可能是pytorch自带的nms底层代码做了优化所以速度会比较快,因此训练的时候不建议用,在使用val.py
进行模型推理的时候可以使用soft-nms
。
一、1️⃣ NMS概念和soft-nms概念
NMS概念:
NMS(Non-maximum suppression)是非极大值抑制, 目的是过滤掉重复的框。 为了保证检测的准确性, 检测网络的输出框一般都比较密集, 对一个物体, 会有多个预测框,NMS就是为了过滤掉这些重复的框, 保留质量最好的那一个框。如下图,设置的阈值可能为0.6,将概率大于0.6的预测框都保留下来,再进行nms非极大值抑制处理,保留与标注框最接近的预测框。
soft-NMS概念:
NMS相比,soft-NM可以理解为是没那么严格的NMS,即使没有达到阈值也会进行第二次判断,例如下面这个图片,上边框都在同一高度,如果使用NMS,可能会把绿色的框删除掉,但是soft-nms会判断它属于另外一匹马,保留该框。
二、2️⃣如何添加soft-nms
2.1 🎓添加soft-nms代码
找到utils文件夹下的general.py文件,在大约880多行的位置找到non_max_suppression函数,如下:
在该函数上方添加如下代码:
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
if CIoU or DIoU or GIoU or EIou:
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 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)
🔥同时找到non_max_suppression函数里的NMS代码,将其注释掉,添加如下代码:
i = soft_nms(boxes, scores, iou_thres)
2.2 ✨添加IOU
在刚刚添加的soft_nms函数里面找到下面这行代码,将你想要使用的Iou修改为True,如果不修改默认的就是普通iou,添加IOU之后就能训练模型了。
三、3️⃣实验结果
➤yolov5基准模型:F1置信度分数为0.71、map@0.5=0.779;
➤使用soft-nms替换掉NMS方法:F1置信度分数为0.81、map@0.5=0.812;虽然map@0.5整体提升了很多,但是该方法对于crazing这个疵点的提升非但没有提升,反而下降了很多。