【pytorch】自己实现精简版YOLOV3【四】，对所得大量预测框进行置信度筛选以及非极大值抑制

接上篇博文：
上一步骤产生大量预测框，其中部分预测框是重复或基本重叠的。需要使用置信度，及框与框间的位置关系（IOU）对所得预测框进行筛选，留下若干最优预测框。如下图所示：

Yolov3中，对预测框的筛选操作如下：
首先，设定置信度阈值，将小于该阈值的预测框全部淘汰。
对剩余的框，通过设定一个阈值（IOU阈值，一般取0.45），判断框之间的距离是否算作接近，并结合置信度得分（yolov3中为物体置信度乘以类别置信度），进行筛选，这种方法又称非极大值抑制（NMS）：
（1）算法输入：包含各框的位置坐标，以及置信度得分
（2）算法步骤：设定IOU阈值，比如0.7。所有边框按照置信度进行排序，取得分最高的边框并保存，使其与其余边框计算IOU，去掉计算结果中大于0.7的结果，保留该边框；再取剩余框中置信度最大的框，与其余边框进行IOU计算，同样舍弃大于阈值（预测框过于接近）的结果，直至所有边框消失。
nms可以用pytorch中torvision库中的nms函数直接实现。举一个例子：

import torch
a=torch.Tensor([[1,1,2,2],[1,1,3.100001,3],[1,1,3.1,3]])
b=torch.Tensor([0.9,0.98,0.980005])
from torchvision.ops import nms
#3个参数为box数组，格式[[x1,y1,x2,y2]],得分tensor[]，及阈值，大于该阈值会被舍弃
ccc=nms(a,b,0.4)
print(ccc)
print(a[ccc])

输出结果为：

tensor([2, 0])
tensor([[1.0000, 1.0000, 3.1000, 3.0000],
        [1.0000, 1.0000, 2.0000, 2.0000]])

非极大值抑制的完整代码为：

```python
import os.path
from typing import Iterator
import numpy as np
import torch
import cv2
import matplotlib.pyplot as plt
import torchvision
from PIL import Image
from torch.utils.data import Dataset, DataLoader, Subset, random_split
import re
from functools import reduce
from torch.utils.tensorboard import SummaryWriter as Writer
from torchvision import transforms, datasets
import torchvision as tv
from torch import nn
import torch.nn.functional as F
import time
import math
def non_max_suppression(prediction, classesNum,conf_thres=0.5, nms_thres=0.4):
    '''
    :param prediction: prediction  [batch_size, 3*inputwidth*inputheight, 4+1+classesNum],
    其中4为x,y,w,h；x和y为网格中心，1为有物体的置信度;inputwidth和inputheight为特征图宽和高的网格数，
    每个特征图的格子中默认预测3个框。
    :param classesNum: 目标检测的类别数目
    :param conf_thres: 有物体的置信度阈值
    :param nms_thres: 非极大值抑制的置信度阈值
    :return:size【n,4】 其中4为box左上角x，y、右下角x，y
    '''
    #首先将x,y,w,h转为左上角和右下角的形式：
    #首先建立一个临时用于存放坐标的坐标：
    boxTemp=torch.FloatTensor(prediction[:,:,0:4].shape)
    boxTemp[:,:,:2]= prediction[:,:,:2] - prediction[:,:,2:4]/2
    boxTemp[:, :,2:4] = prediction[:, :, :2] + prediction[:, :, 2:4] / 2
    prediction[:, :, 0:4]=boxTemp
    '''
    对于prediction  [batch_size, 3*inputwidth*inputheight, 4+1+classesNum]，
    其中[3*inputwidth*inputheight, 4+1+classesNum]代表1张图片，batch_size代表图片的
    数量，本函数要确保函数输出框保持与图片的对应关系，而如下的mask操作，会打乱最外层batchsize维度，
    故用for循环将batchsize维度拆出：
    '''
    #最终输出结果为list，长为bathSize,每个元素为对应的一系列目标框
    outPut=[None]*prediction.shape[0]
    for i,photoInfo in enumerate(prediction):
        '''photoInfo为size[3*inputwidth*inputheight, 4+1+classesNum]'''
    # 首先淘汰所有置信度小于阈值的数据：置信度=类别*类别置信度
        #size(3,1)
        classConf,index=torch.max(photoInfo[:, 5:], dim=1, keepdim=True)
        mask = photoInfo[:, 4]*classConf[:,0] >= conf_thres#*
    #剩余有效数据输出结果为2维：size[n,4+1+classesNum]
        photoInfoR = photoInfo[mask]
        #所对应分出的类别
        classes=index[mask]
        classConf=classConf[mask]
        #识别出类别的集合
        classesSet=classes[-1].unique()
        #保持框数据，置信度与类的对应关系，将三者结合
        #注意，应用NMS时，每次选出置信度最大的框，仍需用到置信度
        #classes要做类别判断，故必须直接放在最后一维（n,4+1+1)
        combine=torch.cat([photoInfoR[:,:4],classConf,classes],-1)
        for classType in classesSet:
            #筛选出同一类型的目标进一步进行nms操作。size[n],返回还是（n2,4+1+1)
            combine=combine[combine[:,-1]==classType]
            boxes=combine[:,:4]
            confList=combine[:,4]
            index = torchvision.ops.nms(boxes, confList, 0.4)
            keep=boxes[index]
            #加入到输出中：
            outPut[i]=keep if outPut[i]==None else torch.cat([outPut[i],keep],dim=0)
    return outPut

#下面进行模拟测试：
k = torch.FloatTensor([x for x in range(60)]).reshape(2, 3, 10)
print(k)
print(non_max_suppression(k, 3))

输出结果为：

[tensor([[ 9.0000,  9.5000, 31.0000, 32.5000],
        [ 4.0000,  4.5000, 16.0000, 17.5000],
        [-1.0000, -0.5000,  1.0000,  2.5000]]), 
 tensor([[24.0000, 24.5000, 76.0000, 77.5000],
        [14.0000, 14.5000, 46.0000, 47.5000]])]