Source-linked AI summary
TPH-YOLOv5: Improved YOLOv5 Based on Transformer Prediction Head for Object Detection on Drone-captured Scenarios
Xingkui Zhu, Shuchang Lyu, Xu Wang, Qi Zhao
TL;DR
Drone-captured object detection is difficult because changing flight altitudes create severe scale variation, while dense scenes and broad coverage complicate object distinction. TPH-YOLOv5 extends YOLOv5 with an extra prediction head, Transformer Prediction Heads, CBAM, and auxiliary training and inference strategies. It achieves 39.18% AP on the VisDrone2021 test-challenge dataset, outperforming DPNetV3 by 1.81%, and ranks fifth in the VisDrone2021 DET challenge.
Problem
Drone-captured scenarios exhibit violent object-scale variation, dense objects with occlusion, and confusing geographic elements, challenging detectors designed for natural scenes.
Method
TPH-YOLOv5 extends YOLOv5 with an additional tiny-object head, Transformer Prediction Heads, CBAM, data augmentation, multiscale testing, model ensembling, and a self-trained classifier.
Results
39.18% AP on the VisDrone2021 test-challenge dataset outperforms DPNetV3 by 1.81%, and TPH-YOLOv5 ranks fifth in the VisDrone2021 DET challenge.
Takeaways & Limitations
TPH-YOLOv5 achieves state-of-the-art performance on VisDrone2021 and is especially effective for object detection in drone-captured scenarios.
Abstract
from arXiv · showhide
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.
1 Beihang University, Beijing, China
The passage provides contact email addresses for the authors.
- Author contact emails are listed as adlith, lyushuchang, sy2002406, and zhaoqi at buaa.edu.cn.
1. Introduction
Drone-captured images challenge object detection through severe scale variation, dense or occluded objects, and confusing geographic backgrounds. TPH-YOLOv5 addresses these issues with additional and transformer-based prediction heads, CBAM, and complementary training and inference strategies, achieving competitive VisDrone2021 results.
- Drone flight-altitude changes produce violent object-scale variation, while dense objects cause occlusion and broad coverage introduces confusing geographic elements.
- TPH-YOLOv5 adds a prediction head for tiny objects, creating four heads for tiny, small, medium, and large objects.
- Transformer Prediction Heads replace the original prediction heads to explore prediction potential with self-attention in challenging dense scenes.
- The method combines data augmentation, multiscale testing, model ensembling, and a self-trained classifier to strengthen drone-image detection.
- CBAM helps the network identify regions of interest and focus on target objects amid confusing geographic information.
- 39.18% AP on the VisDrone2021 test-challenge dataset exceeds DPNetV3 by 1.81%, while TPH-YOLOv5 ranks fifth in the VisDrone2021 DET challenge.
2. Related Work
The section reviews common object-detection architectures and augmentation or ensemble strategies relevant to drone-captured detection. It positions YOLOv5 as a convenient one-stage baseline with CSPDarknet53, PANet, and a YOLO head.
- Data Augmentation: Data augmentation expands training data to improve robustness across images from different environments.The reviewed methods include photometric and geometric distortions, MixUp, CutMix, and Mosaic.
- Data Augmentation: MixUp combines two samples and labels through weighted summation, while CutMix and Mosaic introduce image regions or stitched backgrounds.Mosaic stitches four images and enriches detected-object backgrounds.
- Model Ensemble: Model ensembles reduce neural-network variance by combining predictions from multiple trained models.The section discusses NMS, Soft-NMS, and weighted boxes fusion for combining detection boxes.
- Object Detection: Object detectors commonly include a backbone for feature extraction, a neck between backbone and head, and a head for predicting object locations and categories.The section distinguishes one-stage and two-stage detectors and surveys representative architectures.
- TPH-YOLOv5 Architecture: TPH-YOLOv5's architecture builds on CSPDarknet53, PANet, and transformer prediction heads attached to feature maps from transformer encoder blocks.The architecture figure identifies three transformer encoder blocks at the backbone end and four transformer prediction heads in the neck.
- Object Detection: One-stage detectors predict bounding boxes and classes simultaneously, offering a speed advantage over two-stage detectors but generally lower accuracy.YOLO models are presented as representative one-stage detectors.
3. TPH-YOLOv5
TPH-YOLOv5 adapts YOLOv5 for VisDrone2021 by adding a tiny-object head and transformer-based components, while using attention and inference-time strategies to address scale variation, occlusion, and confusing categories.
- Overview of YOLOv5: YOLOv5 provides the baseline architecture with CSPDarknet53 and SPP as backbone, PANet as neck, and a YOLO detection head.Among YOLOv5 variants, YOLOv5x was selected after showing over a 1.5% AP gap over smaller variants on VisDrone2021.
- Prediction Head for Tiny Objects: An additional high-resolution prediction head targets tiny objects and brings the model to four detection heads for violent object-scale variation.The added head uses a low-level, high-resolution feature map; the extra head increases computation and memory cost.
- Transformer Encoder Block: Transformer encoder blocks replace selected convolutional and CSP bottleneck blocks to capture global and contextual information through self-attention.Each encoder contains multi-head attention and an MLP with residual connections, and performs better on occluded objects in dense VisDrone2021 scenes.
- Transformer Prediction Head: Transformer prediction heads are applied in the head and at the end of the backbone, where low-resolution feature maps reduce computation and memory costs.Some early TPH blocks can be removed when input resolution is enlarged to keep training feasible.
- Convolutional Block Attention Module: CBAM sequentially refines feature maps along channel and spatial dimensions before multiplying the attention map with the input.The paper uses CBAM to focus on useful target objects amid confusing geographical elements in drone images.
- Multi-scale Testing and Model Ensemble: Multi-scale testing and weighted fusion combine predictions from five differently trained models during inference.The pipeline performs multi-scale testing on each model and fuses the five final predictions with WBF.
- Self-trained Classifier: A self-trained ResNet18 classifier addresses low precision in confusing categories such as tricycle and awning-tricycle.The classifier uses 64×64 crops of ground-truth boxes and improves AP by around 0.8%˜1.0%.
4. Experiments
Experiments evaluate TPH-YOLOv5 on VisDrone2021 using challenge and development test sets, reporting mAP and AP50. The study combines implementation choices, multiscale testing, competition results, ablations, and qualitative detections.
- VisDrone2021 evaluation uses testset-challenge and testset-dev, reporting mAP across IoU thresholds [0.5:0.95] and AP50.
- Implementation Details: 65 training epochs with two warm-up epochs are used because the VisDrone2021 training set is relatively small.Training uses Adam, an initial learning rate of 3e-4, cosine decay, 1536-pixel long-side inputs, and batch size 2.
- Data analysis: 622 of 342391 labels are smaller than 3 pixels, and covering these objects with gray squares improves mAP by 0.2.The analysis identifies these very small objects as difficult to recognize.
- Ms-testing: Multiscale testing scales images to three sizes, horizontally flips them, and fuses predictions from six resulting images.The fused prediction is used as the final test result.
- Comparisons with the State-of-the-art: 39.18 AP on testset-challenge ranks TPH-YOLOv5 fifth, 0.25 below the first-place score of 39.43.The reported result comes from an ensemble of five models, while results from four individual models were available.
- Ablation Studies: Adding a tiny-object prediction head increases layers from 607 to 719 and GFLOPs from 219.0 to 259.0, alongside a large mAP improvement.The authors associate the added computation with improved small-object detection.
- Detection results: Qualitative results include large, tiny, dense, and motion-blurred objects, with reported localization performance across these cases.
5. Conclusion
The paper presents TPH-YOLOv5 as a state-of-the-art detector for drone-captured scenarios. Its visualizations show localization of tiny, dense, and motion-blurred objects, while the authors report a VisDrone2021 record.
- TPH-YOLOv5 combines transformer encoder blocks, CBAM, and experienced strategies with YOLOv5 for drone-captured object detection.
- The authors report state-of-the-art performance and a refreshed record on the VisDrone2021 dataset.
- Visualization results show good localization of tiny, dense, and motion-blurred objects.Different categories use bounding boxes with different colors.