yolov5的后处理解析

        由于最近实习项目使用到了yolov5， 发现对yolov5的后处理部分不太熟悉，为防止忘记，这里简单做个记录。

在yolov5里，利用FPN特征金字塔，可以得到三个加强特征层，每一个特征层上每一个特征点存在3个先验框，每个先验框需要预测每一个特征点的回归参数、是否包含物体、以及所包含的物体种类概率。

一、常见预测结果输出类型

一般常见的yolov5模型的输入输出如下面2张图所示，图1为yolov5的三个Predict Head分别输出的示例，图2为3个Head合并后的单结果输出示例。

图1 多结果输出

  图2 单结果输出

假设模型输入为(640,640)的RGB图像，即图中所示输入节点，name：images，shape NCHW=[1,3,640,640]，训练集为coco数据集，具有80个类别，即num_classes=80。

图1中的模型输出分别为(1,255,20,20)，(1,255,40,40)，(1,255,80,80)，其中开头的1表示batch_size，(20,20)、(40,40)、(80,80)分别为3个特征层的形状大小，255=3*(80+4+1)，3表示每个特征点对应的3个先验框，80为one-hot后的各类别概率，4为先验框的回归参数，1为先验框是否包含物体的概率大小。

图2中的25200=3*(20*20+40*40+80*80)，即一共有25200个先验框。

有时模型经过转换后，输出结果为(1,2142000)，其中2142000=25200*85。

二、预测输出结果解码：

由（一）得到的不同预测结果的输出类型，首先先将其reshape一下，变成统一形状（python中对结果列表进行操作即可，注意维度与切分间隔）。

以单结果输出类型(1,25200,85)为例：

其中的85可以拆分成4+1+80。
前4个参数用于判断每一个特征点的回归参数，回归参数调整后可以获得预测框；
第5个参数用于判断每一个特征点是否包含物体；
最后80个参数用于判断每一个特征点所包含的物体种类。

（1）输出结果解码：

主要输出参数有边界框预测的4个offsets：，，，，边界框置信度box_conf 和第i个类别的置信度cls_conf。边界框预测参数需要如下解码公式得到预测框的位置以及宽高。

这里需要注意的是，公式中的解码结果是根据先验框的位置信息进行解码的，如果给定的是先验框的实际位置信息，则解码出来的结果就是实际框的位置信息，这样就不需要针对不同特征层的不同先验框位置进行解码；如果不是，则需要根据不同特征层的先验框位置进行调整，讲个例子，如果公式中给定的、是特征层（20,20）上的坐标网格，那得到的、也是相对于（20,20）上的位置而言的，此时还需要将其映射到模型输入大小上，即乘以其对应的下采样strides。

 第i个类别的置信度cls_conf可以通过获取最大值及其索引得到最大概率max_cls_conf以及相应的索引值max_cls_pred，完成类别预测解码得到种类以及种类置信度。

因此经过解码后可以得到：85(4+1+80) --> 7(x,y,w,h,box_conf,max_cls_conf,max_cls_pred)

（2）非极大值抑制：

经过步骤（1）我们得到了很多的预测框，有些预测框并不准确，因此还需要进行非极大值抑制，防止同一种类框的堆积。

这里借鉴网上的图加深一下影响。

 首先利用置信度进行第一轮的筛选，即筛选出得分满足confidence置信度的预测框，其中框的置信度得分为box_conf*max_cls_con，这样在进行重合框筛选前可以大幅度减少框的数量；

然后进行非极大值抑制，即筛选出一定区域内属于同一种类得分最大的框，

1、获取预测结果中包含的所有种类
2、对种类进行循环，获取该类的全部预测结果
3、根据置信度得分对该种类进行从大到小排序。
4、每次取出得分最大的框，计算其与其它所有预测框的重合程度，重合程度过大的则剔除。

（3）输出结果校正

经过解码以及非极大值抑制后，可以得到模型预测出的相对于输入尺寸大小(640,640)的框，我们需要将预测框校正到实际的图像上。通常实际输入图片比模型输入图像要大，因此需要下采样到模型输入大小，一般采用letterbox操作，即将图像按最小尺度比缩放后在周围补上灰边以填充到模型输入大小，这样如果想得到实际图像上预测框，我们需要对预测结果进一步校正。

1、计算图像大小image_shape与模型输入大小input_shape的最小形状的比值r，得到缩放后的图像形状大小new_shape。

2、计算边界offset以及图像大小与缩放后图像大小的长宽比值scale

3、根据预测框减去边界偏差offset并除以比值得到实际图像上的框的位置与宽高信息。

image_shape=(960,1280)

input_shape=(640,640)

r = min(640/960,640/1280)=min(2/3,0.5)=0.5

new_shape=image_shape*r=(480,640)

offset = (input_shape-new_shape)/2 = ((640,640)-(480,640))/2=(0,80)

scale = 1/r=2

假设模型预测得到了某个框的位置中心点坐标为（100,300），经过上式计算后得到

real = (pred - offset)*scale=[(100,300)-(0,80)]*2=(200,440)

三、相关代码

（1）推荐b导的代码，utils_bbox中实现了torch与numpy的解码。

yolov5-pytorch/utils/utils_bbox.py at main · bubbliiiing/yolov5-pytorch (github.com)

（2）不依赖pytorch，不需要cmake编译，基于官方yolov5训练的模型进行python部署

jndxchengjiabao/yolov5_tensorrt_python: 不使用c++编译，不依赖pytorch，使用tensor和numpy进行推理，导出onnx模型即可使用。 (github.com)

该模型输出为（1,2142000），后处理代码写在trt.py中，使用tensor和numpy实现。 

仅为学习记录，侵删！