【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[