TPH-YOLOv5 | 基于Transformer的YOLOv5小目标检测器

TPH-YOLOv5 | 基于Transformer的YOLOv5小目标检测器 | 四头加注意力

news2024/11/20 4:24:12

在这里插入图片描述
论文地址：https://arxiv.org/pdf/2108.11539.pdf
项目地址：https://github.com/cv516Buaa/tph-yolov5

在无人机捕获的场景中进行对象检测是最近的一项热门任务。由于无人机总是在不同的高度航行，物体尺度变化剧烈，给网络优化带来了负担。此外，高速和低空飞行会在密集的物体上带来运动模糊，这给物体识别带来了很大的挑战。为了解决上述两个问题，我们提出了 TPH-YOLOv5。在 YOLOv5 的基础上，我们增加了一个预测头来检测不同尺度的物体。然后我们用 Transformer Prediction Heads (TPH) 替换原来的预测头，以利用自注意力机制探索预测潜力。我们还集成了卷积块注意模型 (CBAM)，以在具有密集对象的场景中找到注意区域。为了进一步改进我们提出的 TPH-YOLOv5，我们提供了许多有用的策略，例如数据增强、多尺度测试、多模型集成和利用额外的分类器。对数据集 VisDrone2021 的广泛实验表明，TPH-YOLOv5 具有良好的性能，在无人机捕获的场景中具有令人印象深刻的可解释性。在 DET-test-challenge 数据集上，TPH-YOLOv5 的 AP 结果为 39.18%，比之前的 SOTA 方法（DPNetV3）好 1.81%。在 VisDrone Challenge 2021 中，TPHYOLOv5 获得第 5 名，并与第 1 名模型（AP 39.43%）取得了良好的匹配结果。与基线模型（YOLOv5）相比，TPH-YOLOv5 提高了约 7%，令人鼓舞且具有竞争力。

解决的问题

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

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

主要改进

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

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

TPH-YOLOv5网络结构如下：

在这里插入图片描述

TPH模块

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

CBAM模块

在这里插入图片描述

我发现作者公布的代码和图中的代码不一样，所以自己按照上面的图复现了一个，除了检测头以外，完全按照原文内容，这里我们可以参考这篇文章的结构改进自己的模型。因为这些模块我们文件里已经有了，所以我们直接改配置文件就可以了。

# YOLOv5 🚀 by Ultralytics, GPL-3.0 license
# 迪菲赫尔曼 https://blog.csdn.net/weixin_43694096?spm=1000.2115.3001.5343

# Parameters
nc: 80  # number of classes
depth_multiple: 0.33  # model depth multiple
width_multiple: 0.50  # layer channel multiple
anchors:
  - [19,27,  44,40,  38,94]  # P3/8
  - [96,68,  86,152,  180,137]  # P4/16
  - [140,301,  303,264,  238,542]  # P5/32
  - [436,615,  739,380,  925,792]  # P6/64

# YOLOv5 backbone
backbone:
  # [from, number, module, args]
  [[-1, 1, Focus, [64, 3]],       # 0-P1/2
   [-1, 1, Conv, [128, 3, 2]],    # 1-P2/4
   [-1, 3, C3, [128]],
   [-1, 1, Conv, [256, 3, 2]],    # 3-P3/8
   [-1, 9, C3, [256]],
   [-1, 1, Conv, [512, 3, 2]],    # 5-P4/16
   [-1, 9, C3, [512]],
   [-1, 1, Conv, [768, 3, 2]],    # 7-P5/32
   [-1, 1, SPP, [1024, [3, 5, 7]]],
   [-1, 3, C3TR, [1024, False]],  # 9
  ]

# YOLOv5 head
head:
  [[-1, 1, Conv, [768, 1, 1]],  # 10
   [-1, 1, nn.Upsample, [None, 2, 'nearest']], #11
   [[-1, 6], 1, Concat, [1]],   # 12 cat backbone P5
   [-1, 3, C3, [768, False]],   # 13
   [-1, 1, CBAM, [768]],        # 14

   [-1, 1, Conv, [512, 1, 1]],  # 15
   [-1, 1, nn.Upsample, [None, 2, 'nearest']], #16
   [[-1, 4], 1, Concat, [1]],   # 17 cat backbone P4
   [-1, 3, C3, [512, False]],   #  18
   [-1, 1, CBAM, [512]],        # 19

   [-1, 1, Conv, [256, 1, 1]],  # 20
   [-1, 1, nn.Upsample, [None, 2, 'nearest']],  #21
   [[-1, 2], 1, Concat, [1]],   # 22 cat backbone P3
   [-1, 3, C3TR, [256, False]], # 23 (P3/8-small)
   [-1, 1, CBAM, [256]],        # 24

   [-1, 1, Conv, [256, 3, 2]],  # 25
   [[-1, 20], 1, Concat, [1]],  # cat head P4 #26
   [-1, 3, C3TR, [512, False]], # 27 (P4/16-medium)
   [-1, 1, CBAM, [512]],        # 28

   [-1, 1, Conv, [512, 3, 2]],  # 29
   [[-1, 15], 1, Concat, [1]],  # 30 cat head P5
   [-1, 3, C3TR, [768, False]], # 31 (P5/32-large)
   [-1, 1, CBAM, [768]],        # 32

   [-1, 1, Conv, [768, 3, 2]],  # 33
   [[-1, 10], 1, Concat, [1]],  # 34  cat head P6
   [-1, 3, C3TR, [1024, False]],# 35 (P6/64-xlarge)

   [[23, 27, 31, 35], 1, Detect, [nc, anchors]],  # Detect(P3, P4, P5, P6)
  ]