Small object detection in UAV images

Abstract

This paper proposes an improved YOLOv5 algorithm for small object detection in unmanned aerial vehicle (UAV) images. Several modifications are made to enhance the model’s performance, including adding a prediction head to handle large-scale variance, integrating a Channel Feature Fusion with an Involution (CFFI) block to reduce information loss, applying a Convolutional Block Attention Module (CBAM) to focus on important spatial and channel features, and using a C3 structure with a Transformer block (C3TR) to capture contextual information. The proposed method also employs Soft NonMaximum Suppression for improved bounding box scoring in dense scenes. Extensive experiments on the VisDrone2019 dataset demonstrate the effectiveness of these modifications, with the enhanced model outperforming other single-stage detectors and state-of-the-art single-stage detectors like YOLOv8s by a significant margin in terms of mean Average Precision (mAP). Achieving a notable mAP50 of 44.2% and mAP50:95 of 27.3% on the test set. The performance gains are attributed to the integration of attention mechanisms that help the model focus on crucial features for detecting small objects.