TA-YOLO: a lightweight small object detection model based on multi-dimensional trans-attention module for remote sensing images

In recent years, the remarkable performance of convolutional neural networks (CNNs) in computer vision has driven extensive research in deep learning. CNN-based object detection methods can be broadly classified into two categories: two-stage and one-stage detection methods. The former employs a Region Proposal Network (RPN) to generate candidate object boxes, followed by classification and localization. Examples include RCNN [26], Fast R-CNN [27], Faster R-CNN [28], and Mask R-CNN [29], excelling in accuracy but with slower speeds and more parameters.

In contrast, one-stage detection methods directly detect objects in images, eliminating the need for candidate box generation. Notable methods include the YOLO series, MobileNet series [30‐32], ShuffleNet series [33, 34], SSD [35], RetinaNet [36], and EfficientNet [37], known for real-time suitability. The YOLO series has garnered significant attention. YOLOv1 [15] pioneered one-stage detection using Grid Cells for bounding box and class prediction. Subsequent versions improved features, scales, and loss functions. YOLOv4 [18] introduced CSPDarknet53, Spatial Attention Module (SAM), PAN, and CIOU loss. YOLOv5 [19] refined Spatial Pyramid Pooling Fusion (SPPF), diverse training, and balanced loss weights. YOLOv6 [20] focused on efficiency with EfficientRep Backbone and Rep-PAN. YOLOv7 [21] introduced Extended Efficient Layer Aggregation Network (E-ELAN) and Reparameterized Convolution for feature extraction. YOLOv8 [22] advanced with the Compact Context Fusion (C2f) module, Diagonal-Free Loss (DFL) and CIOU Loss, and Task-Aligned Assigner. YOLOv8 is undoubtedly outstanding. However, it still has major limitations in the detection task of small objects.

Small object detection holds significant prominence and poses substantial challenges in the field of computer vision. Remote sensing image data is commonly employed to evaluate the performance of small object detection. Such images often encompass diverse land features, making the targets susceptible to interference from various similar features like color, texture, and shape. Hu et al. [38] utilized a coarse image pyramid and employed twice upsampled input images for detecting small faces. Zheng et al. [39] developed a multi-receptive field convolutional group representation aggregation module to expand the receptive field, thoroughly capturing both the intrinsic features of targets and the semantic relations with the surrounding environment, thereby enhancing detection performance. Liu et al. [40] increased the number of small object training examples by downsizing large objects. D-SSD [41], C-SSD [42], F-SSD [43], and ION [44] focused on constructing suitable context features for small object detection. Furthermore, there are studies utilizing Generative Adversarial Networks (GANs) to generate super-resolution features for small object detection, such as [45] and [46]. However, larger input resolutions and super-resolution methods incur higher computational costs, rendering them unsuitable for lightweight detector designs. Given the nature of small objects, networks often necessitate intricate designs to effectively capture minute features. The aforementioned methods also encounter similar issues of computational complexity. We observe a scarcity of attention devoted to utilizing lightweight object detection frameworks for small object detection tasks. In this paper, our objective is accurate small object detection while maintaining lower computational complexity.

The Attention Mechanism enhances a model’s focus on specific information or regions by adaptively assigning weights to input elements. It finds application in tasks across machine learning, such as natural language processing, computer vision, and speech recognition. Addressing unbalanced foreground-background distribution in object detection, especially for small objects, highlights the significance of attention mechanisms.

The Recurrent Attention Model (RAM) [47] integrates attention with deep neural networks. It iteratively selects and emphasizes image regions, enhancing accuracy and efficiency. SENet [48] introduces attention to convolutional neural networks, learning channel importance with Squeeze-and-Excitation modules. CBAM [49] combines channel and spatial attention for feature enhancement, adapting to diverse image features. Coordinate Attention [50] adds spatial relationship consideration to attention mechanisms, weighting spatial positions for task relevance. Originally in Transformers, the self-attention mechanism computes Query, Key, and Value similarity-based weights for weighted pooling, capturing global dependencies for long-range associations. Inspired by this, we adopt its weight distribution and Multi-Head mechanism to channel and spatial attention extraction.

Transformer is a network based on self-attention mechanisms. The input consists of query and key of dimension \(d_k\), and value of dimension \(d_v\). Multi-head attention is a mechanism that involves employing multiple distinct, learned linear transformations to project the query, key, and value vectors h times onto separate dimensions denoted as \(d_q\), \(d_k\), and \(d_v\). Subsequently, attention operations are conducted in parallel on each of these transformed versions of queries, keys, and values. This orchestrated process results in h sets of output values, each encompassing \(d_v\) dimensions. Transformer architecture further enhances this by amalgamating the outcomes of different attention heads using learnable weights, thereby fostering an adaptive linear aggregation of information. Let q indexes a query element with feature \(x_q\in \mathbb {R}^C\), k indexes a key element with feature \(x_k\in \mathbb {R}^C\), and v indexes a value element with feature \(x_v\in \mathbb {R}^C\), where C is the feature dimension. During the actual operational phase, we simultaneously compute the attention function on a set of queries, packed them into a matrix Q. The keys and values are packed into matrices K and V in the same way. Then they can be expressed as \(X_q\in \mathbb {R}^{d_q}\),\(X_k\in \mathbb {R}^{d_k}\),and \(X_v\in \mathbb {R}^{d_v}\) after being projected h times. Then the multi-head attention feature is calculated bywhere \(\mathbb {C}\) means matrix concatenation. The projections are parameter matrices \(\textit{W}_{i}^Q \in \mathbb {R}^{C\times d_q}\), \(\textit{W}_{i}^K \in \mathbb {R}^{C\times d_k}\), \(\textit{W}_{i}^V \in \mathbb {R}^{C\times d_v}\), and \(\textit{W}_{i}^O \in \mathbb {R}^{hd_v\times C}\). h indexes parallel attention heads and \(d_q=d_k=d_v=C/h\).

The overall structure of TA-YOLO proposed in this paper is shown in Fig. 2. We have closely followed the architectural paradigm of YOLOv8, which similarly comprises three core components: the backbone, neck, and head. In our approach, the notations P2, P3, P4, P5 denote the feature map levels resulting from downsampling the input image by factors of 4, 8, 16, 32 respectively. These levels embody varying degrees of feature map granularity. Much like YOLOv8, our TA-YOLO encompasses a range of models with diverse scales, governed by the predetermined parameters w and n as depicted in the figure. At the termination of the backbone, we have seamlessly integrated the novel MCSTA module to bolster the fundamental feature extractor’s capacity to capture foreground features effectively. Concurrently, this module has been strategically incorporated at the conclusion of each feature level within the neck. This strategic placement serves to further disentangle the foreground and background aspects of the final feature representation.

To aptly address the detection of exceedingly diminutive objects within remote sensing images, we have harnessed the higher-resolution features from the P2 level and introduced a specialized tiny object detection head. Furthermore, we’ve implemented feature fusion between the original P3 and P4 layers in the backbone. This serves to counteract excessive information loss stemming from the iterative upsampling and downsampling processes. By infusing our pioneering MCSTA module and optimizing the network structure, TA-YOLO delivers commendable results across various object sizes. Its efficacy extends beyond medium and large objects, exhibiting superior performance in detecting small and even minuscule objects. This comprehensive design augmentation equips TA-YOLO to excel across a spectrum of object detection scenarios.

The significance of multi-level feature fusion in small object detection cannot be understated. In YOLOv8, object detection relies solely on features from three levels: P3, P4, and P5. However, when dealing with extremely small and densely packed objects within remote sensing images, this approach faces challenges. Features at the P2 level exhibit a higher resolution, offering the network a valuable opportunity to capture intricate texture details present in the image. Building upon this insight, we have introduced a specialized detection head aimed at detecting tiny objects. This innovative addition, depicted by the red dashed box in Fig. 2, enhances our model’s ability to identify minuscule objects.

In the neck architecture, we have integrated the concept of PAN, which operates through a sequence of Down-Up-Down sampling (D-U-D) steps. In YOLOv8, the fusion of features during the second round of D-U-D is limited to the same-level upsampling process. While this helps mitigate information loss to some extent, it remains insufficient for particularly small objects. To address this, we leverage the original features from each level in the backbone. Utilizing residual connections, we compensate for any lost features resulting from repeated upsampling and downsampling processes. Fig. 2 illustrates this approach with red arrows. Notably, the Concat operation at the P5 level amalgamates features processed via SPPF. However, these features remain devoid of repeated upsampling or downsampling. This strategic decision preserves the unique characteristics of these features, preventing excessive fusion and subsequently promoting a more lightweight model architecture.

Through these refined design choices, our model excels in small object detection, benefiting from a comprehensive multi-level feature fusion strategy that harnesses both high-resolution features and precise information propagation.

In the context of small object detection, grappling with the highly imbalanced distribution between foreground and complex background information becomes an unavoidable challenge, particularly pronounced in remote sensing images. This discrepancy presents a formidable obstacle for networks to accurately discern the intricate details of exceedingly small objects within vast-scale images.

To address this issue, we have innovated the MCSTA module. This module serves a dual purpose: firstly, it enhances the extraction of channel trans-attention features via the MCTA submodule; secondly, it proceeds to extract spatial trans-attention features through the MSTA submodule. This sequential process culminates in the fusion of these extracted features with the initial input, yielding a refined and optimized feature representation. The MCSTA module’s architectural blueprint is illustrated in Fig. 3, consisting of the aforementioned MCTA and MSTA sub-modules. To succinctly outline the MCSTA process: consider an input denoted as Z, and the ensuing output as \(Z^{'}\). The MCSTA module can be summarized as follows:where \(\{ \textit{W}_1^{mcta},\textit{W}_2^{msta} \} \in \mathbb {R}^{H\times W\times C}\) represent the weight parameters of the MCTA and MSTA modules, respectively. The design of the residual structure can ensure that our MSCTA module will not interfere with the follow-up results due to attention deviation or confusion during the feature extraction process, and it also plays a supervisory role. By enacting this sophisticated module, we bolster our model’s ability to discern and highlight salient foreground features amidst complex background contexts, thereby elevating its proficiency in detecting small objects within the challenging realm of remote sensing imagery.

The inception of the MCTA and MSTA modules is deeply rooted in the conceptual framework of self-attention, a hallmark of the Transformer architecture. In a parallel trajectory, we retain the fundamental components of query, key, and value from Transformer, harnessing their power to enable information focus and extraction. It’s important to note that while our inspiration draws from Transformer, we have tailored the design to cater to the unique attributes of image data. Conventionally, the Transformer’s architecture divides images into standardized patches, each subsequently encoded with positional information for further processing. In a departure from this approach, we’ve steered clear of employing a block design. Instead, our module operates on the entire image. This modification not only simplifies the model’s complexity but also obviates the need for explicit positional encoding. This streamlined strategy enhances efficiency without compromising performance.

Our module’s adaptation of query, key, and value generation is distinct from Transformer’s methodology. We derive these elements utilizing Ghost convolution [51], an exceptionally lightweight convolutional technique. Ghost convolution leverages linear transformations to eliminate feature redundancies, yielding enriched features with minimal computational overhead. This streamlined procedure, often referred to as a "cheap" operation, omits batch normalization and non-linear activations. This design choice optimally integrates ghost convolutions, emphasizing localized contextual details, culminating in the computation of feature covariances, which in turn generate comprehensive global attention maps.

Intrinsically linked, the MCTA and MSTA modules serve as embodiments of this adapted self-attention principle. Through carefully orchestrated operations, these modules deftly capture nuanced contextual cues, a subject we’ll delve into further in the following elaboration.

The VisDrone dataset [25] was introduced by the AISKYEYE team from Tianjin University’s Machine Learning and Data Mining Laboratory in China. Comprising 288 video clips, this dataset encompasses 261,908 frames and 10,209 still images. It features a diverse range of scenarios and conditions, including different locations, environments, objects, densities, weather, and lighting conditions, captured using multiple drones for varied scenes and missions. The dataset is enriched with 10 categories, containing 6,471 training data, 548 validation data, and 3190 test data samples. Out of these, annotations are available for 1,610 test data samples. In contrast to conventional object detection datasets, the VisDrone dataset often includes hundreds of small objects within a single image, resulting in a substantial total of 2.6 million annotation boxes. Additionally, the dataset provides significant attributes like scene visibility, object class, and occlusion, enhancing its utility for diverse tasks.

The PASCAL VOC dataset [8], originally developed by the Computer Vision group at the University of Oxford, UK, is a well-known and widely employed general-purpose dataset in the field of computer vision. Primarily designed for tasks such as object detection, image segmentation, and image classification, it encompasses 20 common object categories. Spanning multiple years, the VOC dataset was created for the PASCAL VOC challenge. Each year’s dataset contains separate training, validation, and test sets. Notably, data from the years 2007 and 2012 are frequently used. In contemporary practice, the training and validation data from both years are often combined into a unified training set, totaling 16,551 samples. Meanwhile, the test data from the year 2007 serves as the validation or test set, comprising a total of 4952 samples.

The datasets exhibit notable differences in difficulty stemming from their distinct data compositions. Table 1 illustrates this contrast, indicating that 74% of the VisDrone dataset comprises small objects, whereas the PASCAL VOC dataset predominantly contains medium and large objects, accounting for 86%. To offer a visual representation of this divergence, Fig. 4 showcases selected images from both datasets.

The dataset division approach aligns with the methodology outlined in Datasets. In our experimentation, we have adopted several widely recognized evaluation metrics, namely Precision (P), Recall (R), mean Average Precision (mAP), model parameters (Params) and Giga Floating Point Operations Per Second (GFLOPs) [9]. The mean Average Precision (mAP) calculates the average of the average precision (AP) across all categories. This calculation involves utilizing an Intersection over Union (IOU) threshold, with a threshold above which indicating successful detection. We denote the results achieved at an IOU threshold of 0.5 as \(mAP^{50}\). By varying the threshold from 0.5 to 0.95 with an incremental step of 0.05, we compute the average of these values to obtain \(mAP^{50:95}\) [9]. Our implementation is based on PyTorch, with input sample sizes normalized to \(640\times 640\). The training employs the SGD optimizer with an initial learning rate of 0.01, a momentum of 0.937, and a weight decay of 0.0005 to prevent overfitting. We set the batch size to 16 and the maximum epochs to 500, incorporating early stopping with a patience of 50. Remaining hyperparameters are kept consistent with the default YOLOv8 settings. In order to ensure a fair evaluation, all comparative and ablation experiments exclude the use of pre-trained weights during training. Four crucial parameters underpin our TA-YOLO architecture: h, m, n, and w. Here, h represents the heads in the multi-head mechanism, with distinct values 1, 2, 4, 8 considered at each P2, P3, P4, P5 level. \(m=8/h\), its square dictating the number of blocks in the MSTA. The parameter n signifies the number of C2f module stacks, set at 1 for our study. Lastly, w governs the generated model’s size, categorized into five sizes tiny, n, s, o, m, each aligned with w values 0.1875, 0.25, 0.50, 0.652, 0.75, respectively. The loss function undergoes optimization, incorporating DFL Loss in addition to CIOU Loss for enhanced positioning. The comprehensive loss function as follows:where Loss(bce) corresponds to the classification loss function, which encompasses Binary Cross-Entropy (BCE) Loss. On the other hand, Loss(ciou) and Loss(dfl) jointly contribute to form the localization loss function. The parameter \(\lambda \) signifies a weighting factor used to quantify the loss. Notably, specific constants, namely \(\lambda _{1}\), \(\lambda _{2}\), and \(\lambda _{3}\), are assigned values of 0.5, 7.5, and 1.5, respectively, as per YOLOv8’s configuration.

Our primary objective is to design a small object detection model that excels in both speed and accuracy. Given this focus, our experimental comparisons exclude larger models. In the YOLO series, models are typically categorized into five different sizes: n, s, m, l, x. The parameter volume of the YOLOv8-l model has reached 43.7 million (M), so we restrict our analysis to the first three sizes: n, s, m. To offer a comprehensive performance assessment, we introduce two additional small-scale models, resulting in a total of five sizes for our TA-YOLO model: tiny, n, s, o, m.

Our model’s architecture builds upon YOLOv8, prompting a thorough comparison with each size variant of YOLOv8 across the two datasets, as detailed in Tables 2,  3, and  4. Observing these tables reveals notable enhancements in the performance of our TA-YOLO model, achieved with a relatively modest increase in parameters. By comparing the tiny and o-sized models with the n and m sizes-where the parameters are closely matched-it becomes evident that TA-YOLO outperforms YOLOv8 while utilizing fewer parameters. This optimization justifies our augmentation of the tiny and o sizes.

Notably, TA-YOLO boasts superior speed and performance compared to YOLOv8, which is visually demonstrated in Fig. 1. On the PASCAL VOC dataset, our model exhibits improvements of up to 1.3%. However, when evaluated on the validation and test sets of VisDrone, our model achieves remarkable enhancements of up to 6.2% and 4.3% respectively. These findings underscore the exceptional proficiency of our proposed TA-YOLO in small object detection scenarios.

We have extended our comparative analysis beyond YOLOv8 and evaluated our model against other existing methods. As our research emphasizes small object detection, we exclusively employed the VisDrone dataset for these comparisons. The results of these comparative experiments are summarized in Table 5, which highlights the performance of various methods. Notably, our model achieves a compelling combination of minimal parameter quantity and superior accuracy.

Within this context, it’s worth noting that, in addition to the YOLOv7-tiny model, YOLOv7 stands out as having the second lowest parameter count at 36.9M. Consequently, our comparison focuses solely on the YOLOv7-tiny model from the v7 series.

For a more visually intuitive presentation of the comparison results, we’ve crafted Fig. 1, which juxtaposes the two indicators- \(mAP^{50}\) and Params-from Table 5. This graphic depiction offers a concise yet informative representation of how our model outperforms other methods.

The depicted experimental outcomes aptly underscore TA-YOLO’s commendable performance. In our pursuit to delve into our model’s specific performance within the realm of small object detection, we present the Precision-Recall curves for each category on the VisDrone validation set, as vividly illustrated in Fig. 5. These curves manifest the outcomes achieved by setting the IOU threshold to 0.5, subsequently yielding mAP through the integral area under each curve.

Evidently, “car” emerges as the standout performer within this ensemble, capturing the limelight with the highest accuracy. This is particularly noteworthy as “car” claims a substantial share of the dataset. On the contrary, the category with the least representation-namely “Awning-tricycle”-experiences the most challenging detection scenario due to obstructions caused by aerial views and overhanging structures. Consequently, its accuracy is notably lower. The “bicycle” category, with a relatively meager presence, presents further complexities. This challenge arises from occlusions caused by riders and the thin, elongated nature of bicycles at a distance. These factors diminish the availability of discernible and significant features, especially when juxtaposed against the prominence of people within the image. A visual glimpse into these complexities is provided in the example image of Fig. 4.

To provide a more direct and tangible perspective of our method’s efficacy, we present the visual results of detection on selected scenes from the VisDrone dataset in Fig. 6. Evidently, YOLOv8 exhibits noticeable omissions in detecting extremely small objects, particularly in the distant background. The third scene, characterized by a multitude of object categories and instances, presents an intricate challenge exacerbated by occlusions. Notably, YOLOv8 struggles with the detection of pedestrians (indicated by the red bounding box). In contrast, our model noticeably outperforms YOLOv8, demonstrating superior detection capability even for small and occluded objects. This exemplifies our model’s remarkable aptitude for capturing intricate details and overcoming challenges posed by occlusions, leading to enhanced performance in complex scenes.

To validate the effectiveness of our proposed approach in the context of small object detection, we conducted a series of ablation experiments on the VisDrone validation set, utilizing YOLOv8-n as the baseline model. The comprehensive results of these experiments are presented in Table 6. This set of ablation experiments highlights the incremental effect of each module or improvement we introduced, illustrating the diverse degree of performance enhancement achieved. The outcomes of these experiments unequivocally reaffirm the efficacy of our method in elevating network performance across the board.

Furthermore, to explore the impact of integrating the MCSTA module at different network stages, we conducted a series of experiments encompassing five configurations. Initially, we incorporated the module solely in the backbone, then sequentially added it to levels P2 through P5. The experimental results are shown in Table 7. By analyzing the outcomes of these experiments, a clear pattern emerged: the optimal effect is attained by incorporating the module at each stage. Remarkably, this module augmentation incurs a marginal increase in parameters.

To provide a more visual understanding of the efficacy of the MCSTA module, we employed the gradient-based class activation map (Grad-CAM) method [53] to visualize the targeted network layer. Fig. 7 exhibits the feature visualization outcomes for the input detection head within the P3 level. The enhancement afforded by the MCSTA module becomes evident as the model captures a more comprehensive range of foreground information, thereby mitigating background interference while concurrently extracting finer-grained details. This visualization reaffirms the module’s contribution to refining the model’s feature extraction capabilities.