【目标检测】TPH-YOLOv5：基于transformer的改进yolov5的无人机目标检测

简介

最近在使用VisDrone作为目标检测任务的数据集，看到了这个TPH-YOLOv5这个模型在VisDrone2021 testset-challenge数据集上的检测效果排到了第五，mAP达到39.18%。

于是开始阅读它的论文，并跑一跑的它的代码。
论文地址：https://arxiv.org/pdf/2108.11539.pdf
项目地址：https://github.com/cv516Buaa/tph-yolov5
VisDrone数据集下载：https://pan.baidu.com/s/1JzRTeSi_LgdUVhwtbWhA_w?pwd=8888

解决问题

TPH-YOLOv5旨在解决无人机影像中存在的两个问题：

• 因无人机在不同的高度飞行，物体的尺度变化剧烈
• 高速和低空飞行对排列密集的物体带来了运动模糊。

主要改进

TPH-YOLOv5是在YOLOv5的基础上做了下列改进：

• 1、新增了一个检测头来检测更小尺度的物体
• 2、用transformer prediction heads(TPH)替换原来的预测头部
• 3、将CBAM集成到YOLOv5中，帮助网络在大区域覆盖的图像中找到感兴趣的区域。
• 4、其它一系列小tricks

新的检测头

新的检测头不难理解，之前在我的这篇博文【目标检测】YOLOv5针对小目标检测的改进模型/添加帧率检测也提到过这个改进想法。

改进后的网络整体结构图如下：

TPH

作者使用了一个Transformer Encoder来代替一些卷积和CSP结构，将Transformer在视觉中应用，也是目前的主流趋势，Transformer具有独特的注意力机制，效果比原先更好。

CBAM

CBAM(Convolutional Block Attention Module)是作者提出的一种新的设计结构。如图所示，一张特征图被输入到下一个处理单元前，会先并行计算它的通道注意力和空间注意力，然后将其进行融合重塑，这样会让后面的处理单元更加注意到(focus on)有价值的目标区域。

总结，这篇是国人写的论文，论文结构和思路都非常符合中国人的认知习惯，读起来很顺畅。

实战

下面我将使用TPH-YOLOv5对Visdron数据集进行训练。由于代码是根据YOLOv5进行修改的，所以熟悉YOLOv5的读者能够非常轻松的跑通。

值得注意的是，作者提供了两个模型结构，第一个是yolov5l-xs-tph.yaml这个模型结构，并没有用到CBAM，只是在YOLOv5 6.0版本上新增了一个检测头，我估计是消融实验用到的。如果需要跑最好的效果，应该使用yolov5l-xs-tr-cbam-spp-bifpn.yaml这个模型结构。

同时，作者提供了两个预训练模型，之后我会放在文末供读者下载。

我是用Visdron数据集训练100epoch之后，拿网上一段视频来进行检测，和YOLOv5 5.0，6.1版本的结果做对比，效果如下面这个视频所示。

B站Link：https://www.bilibili.com/video/BV17a411u7JD
(前往B站一键三连体验更佳)

可以看到实际效果还是比较明显的，TPH-YOLOv5对于密集人群的识别效果有明显提升。
测试视频我也分享出来：https://pan.baidu.com/s/1jgTonbDYmONkqvLjhLPpRQ?pwd=8888
使用其它模型的测试效果可以@我一下让我康康。

另附测试数据：

注：只是100个epoch的得到的best.pt的测试结果，并未达到最优性能。

代码备份

另附TPH-YOLOv5代码本地备份(包含作者提供的两个预训练权重)：https://pan.baidu.com/s/15mVle5Exghu3jJMFyl9Lyg?pwd=8888