# Pytorch机器学习（八）—— YOLOV5中NMS非极大值抑制与DIOU-NMS等改进

Pytorch机器学习（八）—— YOLOV5中NMS非极大值抑制与DIOU-NMS等改进

# 一、NMS非极大值抑制算法

（这时所有边界框中B1为置信度最高的边界框且没有和其太过相似的边界框——非极大值置信度的边界框被抑制了）

# ; 二、Hard-NMS非极大值代码

"""

socres为N维的tensor，表示每个框的置信度
iou_thres则为上面算法中的IOU阈值

"""
i = torchvision.ops.nms(boxes, scores, iou_thres)  # NMS

def NMS(boxes, scores, iou_thres, GIoU=False, DIoU=False, CIoU=False):
"""

:param boxes:  (Tensor[N, 4])): are expected to be in (x1, y1, x2, y2)
:param scores: (Tensor[N]): scores for each one of the boxes
:param iou_thres: discards all overlapping boxes with IoU > iou_threshold
:return:keep (Tensor): int64 tensor with the indices
of the elements that have been kept
by NMS, sorted in decreasing order of scores
"""
# 按conf从大到小排序
B = torch.argsort(scores, dim=-1, descending=True)
keep = []
while B.numel() > 0:
# 取出置信度最高的
index = B[0]
keep.append(index)
if B.numel() == 1: break
# 计算iou,根据需求可选择GIOU,DIOU,CIOU
iou = bbox_iou(boxes[index, :], boxes[B[1:], :], GIoU=GIoU, DIoU=DIoU, CIoU=CIoU)
# 找到符合阈值的下标
inds = torch.nonzero(iou

def bbox_iou(box1, box2, x1y1x2y2=True, GIoU=False, DIoU=False, CIoU=False, eps=1e-9):
# Returns the IoU of box1 to box2. box1 is 4, box2 is nx4
box2 = box2.T

# Get the coordinates of bounding boxes
if x1y1x2y2:  # x1, y1, x2, y2 = box1
b1_x1, b1_y1, b1_x2, b1_y2 = box1[0], box1[1], box1[2], box1[3]
b2_x1, b2_y1, b2_x2, b2_y2 = box2[0], box2[1], box2[2], box2[3]
else:  # transform from xywh to xyxy
b1_x1, b1_x2 = box1[0] - box1[2] / 2, box1[0] + box1[2] / 2
b1_y1, b1_y2 = box1[1] - box1[3] / 2, box1[1] + box1[3] / 2
b2_x1, b2_x2 = box2[0] - box2[2] / 2, box2[0] + box2[2] / 2
b2_y1, b2_y2 = box2[1] - box2[3] / 2, box2[1] + box2[3] / 2

# Intersection area
inter = (torch.min(b1_x2, b2_x2) - torch.max(b1_x1, b2_x1)).clamp(0) * \
(torch.min(b1_y2, b2_y2) - torch.max(b1_y1, b2_y1)).clamp(0)

# Union Area
w1, h1 = b1_x2 - b1_x1, b1_y2 - b1_y1 + eps
w2, h2 = b2_x2 - b2_x1, b2_y2 - b2_y1 + eps
union = w1 * h1 + w2 * h2 - inter + eps

iou = inter / union
if GIoU or DIoU or CIoU:
cw = torch.max(b1_x2, b2_x2) - torch.min(b1_x1, b2_x1)  # convex (smallest enclosing box) width
ch = torch.max(b1_y2, b2_y2) - torch.min(b1_y1, b2_y1)  # convex height
if CIoU or DIoU:  # 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 distance squared
if DIoU:
return iou - rho2 / c2  # DIoU
elif CIoU:  # https://github.com/Zzh-tju/DIoU-SSD-pytorch/blob/master/utils/box/box_utils.py#L47
v = (4 / math.pi ** 2) * torch.pow(torch.atan(w2 / h2) - torch.atan(w1 / h1), 2)
alpha = v / ((1 + eps) - iou + v)
return iou - (rho2 / c2 + v * alpha)  # CIoU
else:  # GIoU https://arxiv.org/pdf/1902.09630.pdf
c_area = cw * ch + eps  # convex area
return iou - (c_area - union) / c_area  # GIoU
else:
return iou  # IoU

# 四、soft-NMS

si即为置信度，M为当前最大置信度的边界框，bi为其他边界框

[En]

There are mixed opinions on the effect on the Internet, and I haven’t tried it myself, but intuitively, I don’t think the effect will be greatly improved. If you are interested, you can try it yourself.

# ;

Original: https://blog.csdn.net/lzzzzzzm/article/details/120151155
Author: lzzzzzzm
Title: Pytorch机器学习（八）—— YOLOV5中NMS非极大值抑制与DIOU-NMS等改进

(0)