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

简介

最近在使用VisDrone作为目标检测任务的数据集,看到了这个TPH-YOLOv5这个模型在VisDrone2021 testset-challenge数据集上的检测效果排到了第五,mAP达到39.18%。
【目标检测】TPH-YOLOv5:基于transformer的改进yolov5的无人机目标检测_第1张图片
于是开始阅读它的论文,并跑一跑的它的代码。
论文地址: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

新的检测头

【目标检测】TPH-YOLOv5:基于transformer的改进yolov5的无人机目标检测_第2张图片
新的检测头不难理解,之前在我的这篇博文【目标检测】YOLOv5针对小目标检测的改进模型/添加帧率检测也提到过这个改进想法。

改进后的网络整体结构图如下:
【目标检测】TPH-YOLOv5:基于transformer的改进yolov5的无人机目标检测_第3张图片

TPH

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

【目标检测】TPH-YOLOv5:基于transformer的改进yolov5的无人机目标检测_第4张图片

CBAM

【目标检测】TPH-YOLOv5:基于transformer的改进yolov5的无人机目标检测_第5张图片
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版本的结果做对比,效果如下面这个视频所示。

YOLOv5/TPH-YOLOv5检测效果对比测试

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

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

另附测试数据:

算法 [email protected] [email protected]:.95s
yolov5-5.0 34.9% 20.6%
yolov5-6.1 33.1% 18.7%
tph-yolov5 37.4% 21.7%

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

代码备份

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

你可能感兴趣的:(目标检测,目标检测,transformer,深度学习)