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

Xingkui Zhu, Shuchang Lyu+2 more|arXiv (Cornell University)|Aug 26, 2021

Advanced Neural Network Applications56 references122 citations

TL;DR

TPH-YOLOv5 augments YOLOv5 with an extra tiny-object prediction head, Transformer Prediction Heads, and CBAM, plus data-augmentation and ensemble tricks, achieving state-of-the-art on VisDrone2021 test-challenge (AP 39.18%).

ABSTRACT

Object detection on drone-captured scenarios is a recent popular task. As drones always navigate in different altitudes, the object scale varies violently, which burdens the optimization of networks. Moreover, high-speed and low-altitude flight bring in the motion blur on the densely packed objects, which leads to great challenge of object distinction. To solve the two issues mentioned above, we propose TPH-YOLOv5. Based on YOLOv5, we add one more prediction head to detect different-scale objects. Then we replace the original prediction heads with Transformer Prediction Heads (TPH) to explore the prediction potential with self-attention mechanism. We also integrate convolutional block attention model (CBAM) to find attention region on scenarios with dense objects. To achieve more improvement of our proposed TPH-YOLOv5, we provide bags of useful strategies such as data augmentation, multiscale testing, multi-model integration and utilizing extra classifier. Extensive experiments on dataset VisDrone2021 show that TPH-YOLOv5 have good performance with impressive interpretability on drone-captured scenarios. On DET-test-challenge dataset, the AP result of TPH-YOLOv5 are 39.18%, which is better than previous SOTA method (DPNetV3) by 1.81%. On VisDrone Challenge 2021, TPHYOLOv5 wins 5th place and achieves well-matched results with 1st place model (AP 39.43%). Compared to baseline model (YOLOv5), TPH-YOLOv5 improves about 7%, which is encouraging and competitive.

Motivation & Objective

• Address the challenges of drone-captured object detection, including extreme scale variation, high object density, and large scene coverage.
• Enhance YOLOv5 with a dedicated tiny-object head and Transformer-based prediction heads to improve localization and handling of dense scenes.
• Incorporate attention mechanisms and training/inference tricks to boost performance and robustness on drone datasets.

Proposed method

• Add a fourth prediction head to YOLOv5 specifically for tiny objects to handle extreme scale variance.
• Replace original prediction heads with Transformer Prediction Heads (TPH) to utilize self-attention for better localization in crowded scenes.
• Integrate Convolutional Block Attention Module (CBAM) to focus on regions of interest in densely cluttered scenes.
• Apply bag-of-tricks including data augmentation (MixUp, Mosaic), multi-scale testing, and model ensembling to boost accuracy.
• Use a self-trained ResNet18 classifier on cropped object patches to improve misclassified/confused categories and refine final predictions.
• Perform ms-testing by scaling inputs, flipping, and fusing predictions with Weighted Boxes Fusion (WBF) during ensemble.

Experimental results

Research questions

• RQ1 How can Transformer-based prediction heads improve object localization for drone-captured imagery with varying object scales?
• RQ2 What is the impact of adding a tiny-object prediction head and CBAM on detection performance in dense, cluttered drone scenes?
• RQ3 Do data-augmentation, multi-scale testing, and model ensembling significantly improve VisDrone2021 performance, and by how much?
• RQ4 Can a self-trained classifier on cropped patches improve classification accuracy for confusing categories?

Key findings

• TPH-YOLOv5 on VisDrone2021 DET test-dev improves mAP over YOLOv5 baseline and previous ablations.
• Adding a tiny-object head (P2) yields notable AP gains despite higher GFLOPs.
• Transformer encoder blocks reduce network size and GFLOPs while increasing mAP, aiding dense-object detection.
• Model ensemble with ms-testing and WBF achieves higher mAP than single models.
• A self-trained classifier provides about 0.8–1.0% AP improvement on final results.
• On VisDrone2021 test-challenge, TPH-YOLOv5 ensemble achieves 39.18% AP, outperforming prior SOTA DPNetV3 by 1.81% (Table 1).