TPH-YOLOv5 论文笔记

TPH-YOLOv5: Improved YOLOv5 Based on Transformer Prediction Head for Object Detection on Drone-captured Scenarios

TPH-YOLOv5 2021 ICCV workshop

论文链接: https://arxiv.org/abs/2108.11539

一、 Problem Statement

利用和改进YOLOv5进行无人机目标识别。无人机目标识别通常有以下几个问题:

1. 目标大小变化较大
2. 运动模糊且目标密度大，目标覆盖重叠多
3. 杂乱的地理元素，也就是目标的背景信息多

二、 Direction

基于YOLOv5，作者做了以下改进:

1. 增加一个Prediction head。
2. 使用Transformer encoder blocks替换一些convolutional blocks和CSP bottleneck blocks。
3. 增加了CBAM block。
4. Multi-scale testing 和 self-trained classifier。

三、 Method

先来看一下整体的网络框架:

1. Prediction head for tiny objects

对于小目标检测，作者添加了一个Prediction head，如上图的head 1。这样总共四个Prediction head可以减轻目标大小变化较大的影响。所添加的那个Prediction head的输入是low-level和high-resolution的特征图，对小目标比较敏感。尽管添加了一个prediction head会带来速度影响，但是精度却提高较多。

2. Transformer encoder block

用transformer encoder block 替换一些convolutional block和CSP bottleneck blocks。作者认为transformer encoder block可以捕获全局信息和充足的背景信息，而且高密度遮挡目标上的表现也更好。

作者在prediction head上和backbone的末尾使用了transformer。因为网络的末端的特征图分辨率较小，可以降低计算和存储消耗。

3. CBAM

在无人机拍摄的图像上，大的覆盖区域总是包含令人困惑的地理元素。使用CBAM可以提取注意区域，帮助TPH-YOLOv5抵抗混乱的信息，并专注于有用的目标对象。

4. Multi-scale testing and self-trained classifier

作者发现TPH-YOLOv5有较好的定位能力，但是分类能力较弱。因此提出了一个额外的self-trained classifier。首先会通过裁剪ground-truth bounding boxes和resize到64x64来构建训练集。然后使用ResNet18，进行训练。

每个模块对应的提升效果如下:

四、 Conclusion

对于改进YOLOv5提供了tricks，可以借鉴参考。

Reference