Improved YOLOv5 for Aerial Images Based on Attention Mechanism

Abstract

Object detection based on unmanned aerial vehicle (UAV) platforms is essential for both engineering and research. Complex scale problems in UAV application scenarios require strong regression localization capabilities from target detection algorithms. Nonetheless, due to the constraints of UAV platform, it is difficult to increase accuracy by deepening the network. Therefore, this paper presents an improved YOLOv5 with an attention mechanism, consisting a Convolution-Swin Transformer Block (CSTB) utilizing Swin Transformer as well as a Convolution-block Attention Module (CBAM) to improve network positioning accuracy. In addition, this paper incorporates Bidirectional Feature Pyramid Network (BiFPN), Spatial Pyramid Pooling-Fast (SPPF) and some network components to increase the average precision while maintaining the limited size of the model. Experiments on Visdrone2019 dataset show that the proposed approach can raise the mean Average Precision (mAP) by 5.4% compared to YOLOv5, with only 18% increase in model size.