【python】实现NMS

NMS

NMS非极大值抑制，目标：消除多余的框，只保留最佳物体检测的位置

步骤：

1. 对所有检测框按置信度进行排序
2. 找出分数最高的检测框$B_{max}$
3. 遍历其他框，计算当前框与$B_{max}$的重叠面积
4. 如果重叠面积大于某个设定的阈值，则删除
5. 继续对未处理的检测框进行排序
6. 重复

检测模型的输出为：(batch_size, all_boxes, 4+1+num_classes)

对于yolo的输出结果具体为：
batch_size, x,y,w,h, 存在物体的概率，类别1的概率…类别n的概率

对于SSD的输出结果具体为：
batch_size, x,y,w,h, 存在背景的概率，类别1的概率…类别n的概率

实现NMS

'''对于yolo的nms代码'''

def nms(boxes, num_classes, conf_thres=0.5, nms_thres=0.4):
 #boxes: 检测的boxes及scores，维度(batch_size, all_boxes, 4+1+num_classes)
 #num_classes: 类别的数量
 #conf_thres: 预测框的阈值, nms_thres: 极大值抑制的阈值


 #取batch size
 bs = np.shape(boxes)[0]

 #将预测框（中心宽高）转成左上角右下角的形式
 shape_boxes = np.zeros_like(boxes[:,:,:4])
 shape_boxes[:,:,0] = boxes[:,:,0] - boxes[:,:,2]/2
 shape_boxes[:,:,1] = boxes[:,:,1] - boxes[:,:,3]/2
 shape_boxes[:,:,2] = boxes[:,:,0] + boxes[:,:,2]/2
 shape_boxes[:,:,3] = boxes[:,:,1] + boxes[:,:,3]/2

 boxes[:,:,:4] = shape_boxes
 output = []
 
 #遍历每一个batch
 for i in range(bs):
  # prediction shape (num_boxes, 4+1+num_classes)
  prediction = boxes[i]

  #获取存在物体的概率
  score = prediction[:, 4]

  #利用mask筛选出一张图像里面置信度大于设定值的预测框
  mask = score > conf_thres
  detection = prediction[mask]

  #找到这些预测框属于什么种类
  #取prediction (num_boxes, 4+1+num_classes)最后一个值
  # 对其求max，找出属于某种类的概率
  class_conf = np.expand_dims(np.max(detections[:, 5:], axis=-1), axis=-1)
  #对其求argmax，找出属于哪个种类argmax返回最大值的索引
  class_pred = np.expand_dims(np.argmax(detections[:, 5:], axis=-1), axis=-1)
  
  #堆叠，得到的维度(num_boxes, 4+1+2) : 2代表所属类别的置信度和种类
  detections = np.concatenate([detections[:, :5], class_conf, class_pred], axis=-1)

  #对种类进行一个去重的操作
  unique_class = np.unique(detection[:, -1])
  if(len(unique_class)==0):
   continue
  
  #存放nms的结果
  best_box = []
  
  #遍历每个类别
  for c in unique_class:
   # 获取该类别的预测框
   cls_mask = detections[:, -1] == c 
   detection = detections[cls_mask]

   #对得分进行从大到小排序
   scores = detection[:, 4]
   arg_sort = np.argsort(scores)[::-1] #[::-1]取反，从大到小排序
   detection = detection[arg_sort]

   #
   while len(detection) != 0:
    #首先取出最大的框，保存到结果里
    best_box.append(detection[0])

    if len(detection)==1:
     break

    #计算最大框和其他框的重合程度，重合程度大的剔除
    ious = IoU(best_box[-1], detection[1:])
    detection = detection[1:][ious<nms_thres]

  output.append(best_box)
 return np.array(output)




def IoU(b1, b2):
 #计算两个框的重合程度
 b1_x1, b1_y1, b1_x2, b1_y2 = b1[0], b1[1], b1[2], b1[3]
 b2_x1, b2_y1, b2_x2, b2_y2 = b1[:, 0], b1[:, 1], b2[:, 2], b2[:, 3]
  
 #求交集
 inter_rect_x1 = np.maximum(b1_x1, b2_x1) 
 inter_rect_y1 = np.maximum(b1_y1, b2_y1)
 inter_rect_x2 = np.maximum(b1_x2, b2_x2)
 inter_rect_y2 = np.maximum(b1_y2, b2_y2)

 inter_area = np.maximum(inter_rect_x2 - inter_rect_x1, 0) * \
       np.maximum(inter_rect_y2 - inter_rect_y1, 0)

 #求并集
 area_b1 = (b1_x2 - b1_x1) * (b1_y2 - b1_y1)
 area_b2 = (b2_x2 - b2_x1) * (b2_y2 - b2_y1)

 #求iou值
 iou = inter_area / np.maximum(area_b1 + area_b2 - inter_area, 1e-6)
 return iou
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102

Soft-NMS（软NMS）

NMS能解决大部分重叠问题，但当两个目标挨得比较近时，如果对于另一个目标的预测框和当前目标的预测框重叠较大的情况，NMS会直接剔除掉另一个目标的预测框，如下图：

最后只保留红色框，绿色框会被删除。

因此，Soft-NMS的目标是不要直接删除所有IOU大于阈值的框，而是降低置信度。它的算法流程如下图：

Soft-NMS 有两种形式：

1. 线性加权：
   S i = { S i IOU(M, bi)=t S_i=
   {SiIOU(M, bi)<tSi∗(1−IOU(M,bi))IOU(M,bi)>=t" role="presentation" style="position: relative;">{SiSi∗(1−IOU(M,bi))IOU(M, bi)IOU(M,bi)>=t