Enhanced Detection and Recognition of Road Objects in Infrared Imaging Using Multi-Scale Self-Attention
Abstract
:1. Introduction
- Addressing the challenges of infrared small target detection, which is susceptible to background interference and typically suffers from low accuracy, this paper builds upon YOLOv8 by harnessing the efficient feature transmission capabilities of attention mechanisms. Through the use of self-attention across multiple scales, the model has been enhanced. To prevent model bloat, pruning techniques were employed to reduce the model’s parameter count. Experiments conducted on both proprietary and public datasets demonstrate that our model exhibits robustness and state-of-the-art performance.
- Improvements across multiple scales: To address information loss in the traditional YOLO algorithm via the FPN structure, this work incorporates a model modification using a Gather-and-Distribute Mechanism constructed from self-attention. This significantly enhances the model’s ability to detect small objects. Additionally, the original detection head was improved by introducing a DyHead structure formed from multi-head self-attention, enabling the model to rapidly extract crucial information from the feature layers.Among them, * represents the operation between the feature map and the attention function
- Model light-weighting operations: To facilitate better deployment of our model, we conducted further model pruning operations on the developed model. After pruning, the size of the model and the number of parameters were significantly reduced, while maintaining an acceptable level of accuracy loss. This process involves systematically removing less important parameters within the neural network, which not only reduces the computational burden but also improves the efficiency of the model. This optimization ensures that the model remains robust and agile, capable of performing effectively without compromising the essential detection capabilities.
- Related experiments: The enhanced model was compared against different datasets and baseline models. Feature map visualization was employed to illustrate the improvements in model accuracy. The experimental results indicate that the proposed improved model has strong generalization capabilities. Our method achieves a 3.1% improvement in AP and a 2.5% increase in mAP on the VisDrone2019 dataset.
2. Related Work
2.1. Introduction to YOLO Algorithm
- Unlike traditional sliding window and region proposal network-based detection algorithms, YOLO transforms the object detection task into a regression problem, processing the entire image in one go. This significantly enhances the detection speed. According to recent research data, YOLO can process images at 45 frames per second, while the fastest algorithm in the R-CNN series [21], Fast R-CNN [22], processes only 7 frames per second. This high efficiency makes YOLO particularly suitable for scenarios where real-time response is crucial [23].
- YOLO employs a global optimization strategy, considering the overall information of the image rather than just focusing on local features [26]. This capability enables YOLO to handle some challenges that traditional detection algorithms struggle with, such as detecting small or overlapping objects, with higher accuracy [27].
- The network structure of YOLO is simple, making it easy to integrate with other neural network structures [28]. For example, additional convolutional layers can be added to YOLO for feature extraction, or fully connected layers can be added for classification. This ease of integration allows YOLO to perform multiple tasks simultaneously in composite applications, such as pedestrian and vehicle detection in autonomous driving, while maintaining high processing speed and accuracy.
2.2. Introduction to Attention Mechanisms
- Linear transformations: Each head in the multi-head attention mechanism performs separate linear transformations of the input X to create different sets of queries, keys, and values. This is performed using parameter matrices , , and , which are learned during training:These matrices , , and are unique to each head, allowing the model to learn different aspects of the data in each subspace.
- Scaled Dot-product Attention: Each head computes the attention scores using the Scaled Dot-product Attention formula:Here, the scaling factor normalizes the dot products, as mentioned earlier, to avoid extremely large values that could destabilize the softmax function.
- Concatenation of heads: After each head computes its output, the outputs are concatenated along the feature dimension:
- Output transformation: Finally, the concatenated output from all the heads is once again linearly transformed by to yield the final output of the multi-head attention layer. This output can then be used in subsequent layers of the model.
- Input the feature map: The input feature map has dimensions H × W × C.
- Spatial feature Squeeze: Compress the spatial features of the input feature map by applying global average pooling across the spatial dimensions, resulting in a feature map of size 1 × 1 × C.
- Channel feature learning: Learn channel-specific features from the compressed feature map by passing it through a fully connected (FC) layer. This operation produces a feature map with channel attention, maintaining the dimensions 1 × 1 × C
- Channel attention application: Multiply the channel-attended feature map 1 × 1 × C with the original input feature map H × W × C on a channel-by-channel basis to apply the learned attention weights. The final output is a feature map with enhanced channel attention. In this process, the FC layer predicts the importance of each channel, applying different weights based on the importance of each channel.
3. Method
3.1. DyHead Detection Head Improvement
- Scale-aware attention is implemented at the level dimension, providing scale-sensitive attention. As different levels of feature maps correspond to different target scales, the introduction of this mechanism enhances the model’s ability to perceive scale variations.
- Spatial-aware attention is implemented at the spatial dimension, providing spatial-sensitive attention. Changes in spatial positioning correspond to geometric transformations of the detection targets, and the introduction of this mechanism enhances the model’s ability to perceive spatial positions.
- Task-aware attention is implemented at the channel dimension, providing task-sensitive attention. Different channels correspond to different detection tasks, and the introduction of this mechanism enhances the model’s ability to perceive and differentiate between various tasks.
3.2. Gather-and-Distribute Mechanism Improvement
3.2.1. Low-GD Branch
- Low-FAM (Low-level Feature Alignment Module). As depicted in the diagram, B4 serves as the reference for alignment. Larger feature maps such as B2 and B3 are downsampled through average pooling, while smaller feature maps like B5 are upsampled using bilinear interpolation to standardize the size of the feature maps. The aligned feature maps are then concatenated to form the merged feature . This is followed by processing through a Transformer module, which minimizes computational complexity.
- Low-IFM (Low-level Information Fusion Module). The design of this module includes a Convolution (Conv) module, RepBlock module, and a Split operation. The aligned and concatenated features are input into the RepBlock module to obtain the fused features . The Conv module adjusts the channel dimensions to accommodate different model sizes. is then split along the channel dimension into and , which are used for subsequent feature fusion with different layers. The equations are as follows:
- Inject (Information Distribution Module). The distribution of information is facilitated using a self-attention mechanism, as illustrated in Figure 9. This approach ensures that each layer receives the most relevant and contextually appropriate information, enhancing the model’s overall detection capability.
3.2.2. High-GD Branch
3.3. Model Pruning
- Pretrain: Pretraining involves regular training procedures to develop a model that is subsequently subjected to constrained training.
- Constrained training: To promote structural sparsity within the YOLOv8 infrared small target detection model, L1 regularization is applied to the coefficients of the Batch Normalization (BN) layers. Constrained training facilitates easier pruning by allowing the model to identify channels or weight coefficients that are less important. These elements are then targeted in the pruning process to achieve model compression. Direct pruning without this preliminary step might lead to significant issues such as a drastic reduction in model accuracy and uneven pruning. Therefore, sparse training before pruning enables more precise identification of expendable channels or weights, thus preventing such problems.
- Pruning: After sparse training, pruning is performed according to a predefined rate to create a compact model that requires less storage capacity.
- Model fine-tuning: The primary goal of fine-tuning is to effectively recover the accuracy that might be compromised due to significant losses caused by channel pruning.
4. Experiment
4.1. Model Evaluation Metrics
- Precision (P): This is the ratio of correctly predicted positive observations to the total predicted positives. It is calculated as:
- Recall (R): This is the ratio of correctly predicted positive observations to the all observations in actual class. It is defined as:
- Mean average precision (mAP): For multiple classes, mAP is the mean of the average precision calculated for each class. Average precision (AP) for a single class is the area under the precision–recall curve for that class. The mAP across all classes is calculated by averaging the APs:
4.2. Dataset Creation and Experiments
4.3. Ablation Study
4.4. Target Heatmaps and Detection Effects
5. Conclusions
Author Contributions
Funding
Institutional Review Board Statement
Informed Consent Statement
Data Availability Statement
Conflicts of Interest
References
- Peri, S.D.B.; Palaniswamy, S. A Novel Approach To Detect and Track Small Animals using YOLOv8 and DeepSORT. In Proceedings of the 2023 4th IEEE Global Conference for Advancement in Technology (GCAT), Bangalore, India, 6–8 October 2023; pp. 1–6. [Google Scholar] [CrossRef]
- Wu, Z.; Yu, H.; Zhang, L.; Sui, Y. AMB: Automatically Matches Boxes Module for One-Stage Object Detection. In Proceedings of the 2023 IEEE International Conference on Image Processing and Computer Applications (ICIPCA), Changchun, China, 11–13 August 2023; pp. 1516–1522. [Google Scholar] [CrossRef]
- Yang, J.; Ni, Q.; Luo, G.; Cheng, Q.; Oukhellou, L.; Han, S. A Trustworthy Internet of Vehicles: The DAO to Safe, Secure, and Collaborative Autonomous Driving. IEEE Trans. Intell. Veh. 2023, 8, 4678–4681. [Google Scholar] [CrossRef]
- Du, J.; Lu, H.; Zhang, L.; Hu, M.; Shen, X. Infrared Small Target Detection and Tracking Method Suitable for Different Scenes. In Proceedings of the 2020 IEEE 9th Joint International Information Technology and Artificial Intelligence Conference (ITAIC), Chongqing, China, 11–13 December 2020; Volume 9, pp. 664–668. [Google Scholar] [CrossRef]
- Lin, T.Y.; Dollár, P.; Girshick, R.; He, K.; Hariharan, B.; Belongie, S. Feature Pyramid Networks for Object Detection. arXiv 2017, arXiv:1612.03144. [Google Scholar]
- Zhang, T.; Gai, K.; Bai, H. Multiscale Image Deblurring Network Using Dual Attention Mechanism. In Proceedings of the 2022 16th IEEE International Conference on Signal Processing (ICSP), Beijing, China, 21–24 October 2022; pp. 85–89. [Google Scholar] [CrossRef]
- Lin, T.Y.; Goyal, P.; Girshick, R.; He, K.; Dollár, P. Focal Loss for Dense Object Detection. arXiv 2018, arXiv:1708.02002. [Google Scholar]
- Liu, S.; Qi, L.; Qin, H.; Shi, J.; Jia, J. Path Aggregation Network for Instance Segmentation. arXiv 2018, arXiv:1803.01534. [Google Scholar]
- Yan, L.; Li, K.; Gao, R.; Wang, C.; Xiong, N. An Intelligent Weighted Object Detector for Feature Extraction to Enrich Global Image Information. Appl. Sci. 2022, 12, 7825. [Google Scholar] [CrossRef]
- Feng, B.; Jin, Y.; Yin, Z.; Liu, Y.; Wang, X.; Zhao, Y. Infrared Image Recognition and Classification of Typical Electrical Equipment in Substation Based on YOLOv5. In Proceedings of the 2023 2nd Asian Conference on Frontiers of Power and Energy (ACFPE), Chengdu, China, 20–22 October 2023; pp. 140–144. [Google Scholar] [CrossRef]
- Zhang, X.; Zhu, X. Vehicle Detection in the Aerial Infrared Images via an Improved Yolov3 Network. In Proceedings of the 2019 IEEE 4th International Conference on Signal and Image Processing (ICSIP), Wuxi, China, 19–21 July 2019; pp. 372–376. [Google Scholar] [CrossRef]
- Eysa, R.; Hamdulla, A. Issues on Infrared Dim Small Target Detection and Tracking. In Proceedings of the 2019 International Conference on Smart Grid and Electrical Automation (ICSGEA), Xiangtan, China, 10–11 August 2019; pp. 452–456. [Google Scholar] [CrossRef]
- Wei, H.; Tan, Y.; Lin, J. Robust Infrared Small Target Detection via Temporal Low-Rank and Sparse Representation. In Proceedings of the 2016 3rd International Conference on Information Science and Control Engineering (ICISCE), Beijing, China, 8–10 July 2016; pp. 583–587. [Google Scholar] [CrossRef]
- Redmon, J.; Divvala, S.; Girshick, R.; Farhadi, A. You Only Look Once: Unified, Real-Time Object Detection. arXiv 2016, arXiv:1506.02640. [Google Scholar]
- Redmon, J.; Farhadi, A. YOLO9000: Better, Faster, Stronger. arXiv 2016, arXiv:1612.08242. [Google Scholar] [CrossRef]
- Redmon, J.; Farhadi, A. YOLOv3: An Incremental Improvement. arXiv 2018, arXiv:1804.02767. [Google Scholar]
- Bochkovskiy, A.; Wang, C.Y.; Liao, H.Y.M. YOLOv4: Optimal Speed and Accuracy of Object Detection. arXiv 2020, arXiv:2004.10934. [Google Scholar]
- Gai, R.; Li, M.; Chen, N. Cherry detection algorithm based on improved YOLOv5s network. In Proceedings of the 2021 IEEE 23rd International Conference on High Performance Computing & Communications; 7th International Conference on Data Science & Systems; 19th International Conference on Smart City; 7th International Conference on Dependability in Sensor, Cloud & Big Data Systems & Application (HPCC/DSS/SmartCity/DependSys), Haikou, China, 20–22 December 2021; pp. 2097–2103. [Google Scholar] [CrossRef]
- Zhang, H.; Guo, H.; Ding, K.; Liu, J.; Chen, W. Complex Small Target Image Recognition Algorithm Based on Data Enhancement in YOLOv7. In Proceedings of the 2023 IEEE 6th International Conference on Pattern Recognition and Artificial Intelligence (PRAI), Haikou, China, 18–20 August 2023; pp. 470–473. [Google Scholar] [CrossRef]
- Wang, C.Y.; Yeh, I.H.; Liao, H.Y.M. YOLOv9: Learning What You Want to Learn Using Programmable Gradient Information. arXiv 2024, arXiv:2402.13616. [Google Scholar]
- Liu, Z.; Deng, Y.; Ma, F.; Du, J.; Xiong, C.; Hu, M.; Zhang, L.; Ji, X. Target detection and tracking algorithm based on improved Mask RCNN and LMB. In Proceedings of the 2021 International Conference on Control, Automation and Information Sciences (ICCAIS), Xi’an, China, 14–17 October 2021; pp. 1037–1041. [Google Scholar] [CrossRef]
- Duth, S.; Vedavathi, S.; Roshan, S. Herbal Leaf Classification using RCNN, Fast RCNN, Faster RCNN. In Proceedings of the 2023 7th International Conference on Computing, Communication, Control and Automation (ICCUBEA), Pune, India, 18–19 August 2023; pp. 1–8. [Google Scholar] [CrossRef]
- Luo, Q.; Zhang, Z.; Yang, C.; Lin, J. An Improved Soft-CBAM-YoloV5 Algorithm for Fruits and Vegetables Detection and Counting. In Proceedings of the 2023 IEEE 6th International Conference on Pattern Recognition and Artificial Intelligence (PRAI), Haikou, China, 18–20 August 2023; pp. 187–192. [Google Scholar] [CrossRef]
- Pandey, S.; Chen, K.F.; Dam, E.B. Comprehensive Multimodal Segmentation in Medical Imaging: Combining YOLOv8 with SAM and HQ-SAM Models. In Proceedings of the 2023 IEEE/CVF International Conference on Computer Vision Workshops (ICCVW), Paris, France, 2–6 October 2023; pp. 2584–2590. [Google Scholar] [CrossRef]
- Li, G.; Zhu, D. Research on road defect detection based on improved YOLOv8. In Proceedings of the 2023 IEEE 11th Joint International Information Technology and Artificial Intelligence Conference (ITAIC), Chongqing, China, 8–10 December 2023; pp. 143–146. [Google Scholar] [CrossRef]
- Gunawan, F.; Hwang, C.L.; Cheng, Z.E. ROI-YOLOv8-Based Far-Distance Face-Recognition. In Proceedings of the 2023 International Conference on Advanced Robotics and Intelligent Systems (ARIS), Taipei, Taiwan, 30 August–1 September 2023; pp. 1–6. [Google Scholar] [CrossRef]
- Yan, S.; Fu, Y.; Zhang, W.; Yang, W.; Yu, R.; Zhang, F. Multi-Target Instance Segmentation and Tracking Using YOLOV8 and BoT-SORT for Video SAR. In Proceedings of the 2023 5th International Conference on Electronic Engineering and Informatics (EEI), Wuhan, China, 30 June–2 July 2023; pp. 506–510. [Google Scholar] [CrossRef]
- Samaniego, L.A.; Peruda, S.R.; Brucal, S.G.E.; Yong, E.D.; De Jesus, L.C.M. Image Processing Model for Classification of Stages of Freshness of Bangus using YOLOv8 Algorithm. In Proceedings of the 2023 IEEE 12th Global Conference on Consumer Electronics (GCCE), Nara, Japan, 10–13 October 2023; pp. 401–403. [Google Scholar] [CrossRef]
- Zhang, Y.F.; Sun, T.; Ma, S. Improved Generative Adversarial Network with Hybrid Attention Mechanism for Path Planning. In Proceedings of the 2023 42nd Chinese Control Conference (CCC), Tianjin, China, 24–26 July 2023; pp. 5933–5938. [Google Scholar] [CrossRef]
- Vaswani, A.; Shazeer, N.; Parmar, N.; Uszkoreit, J.; Jones, L.; Gomez, A.N.; Kaiser, L.; Polosukhin, I. Attention Is All You Need. arXiv 2023, arXiv:1706.03762. [Google Scholar]
- He, Y.; Huang, H.; Li, Q.; Zhang, G. Research on attention-based multiscale information fusion with the real-time instance segmentation method. In Proceedings of the 2022 2nd International Conference on Algorithms, High Performance Computing and Artificial Intelligence (AHPCAI), Guangzhou, China, 21–23 October 2022; pp. 568–572. [Google Scholar] [CrossRef]
- Shi, Y.; Hidaka, A. Attention-YOLOX: Improvement in on-Road Object Detection by Introducing Attention Mechanisms to YOLOX. In Proceedings of the 2022 International Symposium on Computing and Artificial Intelligence (ISCAI), Beijing, China, 16–18 December 2022; pp. 5–14. [Google Scholar] [CrossRef]
- Yang, J.; Gu, Y.; Sun, Z.; Cui, Z. A Small Infrared Target Detection Method Using Adaptive Local Contrast Measurement. In Proceedings of the 2019 IEEE International Instrumentation and Measurement Technology Conference (I2MTC), Auckland, New Zealand, 20–23 May 2019; pp. 1–6. [Google Scholar] [CrossRef]
- Talaat, F.M.; ZainEldin, H. An improved fire detection approach based on YOLO-v8 for smart cities. Neural Comput. Appl. 2023, 35, 20939–20954. [Google Scholar] [CrossRef]
- Dai, X.; Chen, Y.; Xiao, B.; Chen, D.; Liu, M.; Yuan, L.; Zhang, L. Dynamic Head: Unifying Object Detection Heads with Attentions. arXiv 2021, arXiv:2106.08322. [Google Scholar] [CrossRef]
- Umanandhini, D.; Devi, M.S.; Beulah Jabaseeli, N.; Sridevi, S. Batch Normalization based Convolutional Block YOLOv3 Real Time Object Detection of Moving Images with Backdrop Adjustment. In Proceedings of the 2023 9th International Conference on Smart Computing and Communications (ICSCC), Kochi, India, 17–19 August 2023; pp. 25–29. [Google Scholar] [CrossRef]
- Rooban, S.; S, I.J.; Manimegalai, R.; Eshwar, I.V.S.; Mageswari, R.U. Simulation of Pick and Place Robotic Arm using Coppeliasim. In Proceedings of the 2022 6th International Conference on Computing Methodologies and Communication (ICCMC), Erode, India, 29–31 March 2022; pp. 600–606. [Google Scholar] [CrossRef]
- Cao, X.; Yao, R.; Yao, P. Targets Recognition Based on Feature Extraction with Small Convolutional Networks for Ultrasonic NDT. In Proceedings of the 2023 International Conference on Intelligent Management and Software Engineering (IMSE), Rome, Italy, 8–10 September 2023; pp. 108–111. [Google Scholar] [CrossRef]
- Singh, R.; Singh, D. Quality Inspection with the support of computer vision techniques. In Proceedings of the 2022 International Interdisciplinary Humanitarian Conference for Sustainability (IIHC), Bengaluru, India, 18–19 November 2022; pp. 1584–1588. [Google Scholar] [CrossRef]
- Selvaraju, R.R.; Cogswell, M.; Das, A.; Vedantam, R.; Parikh, D.; Batra, D. Grad-CAM: Visual Explanations from Deep Networks via Gradient-based Localization. Int. J. Comput. Vis. 2020, 128, 336–359. [Google Scholar] [CrossRef]
- Du, D.; Zhu, P.; Wen, L.; Bian, X.; Ling, H.; Hu, Q.; Zheng, J.; Peng, T.; Wang, X.; Zhang, Y.; et al. VisDrone-SOT2019: The Vision Meets Drone Single Object Tracking Challenge Results. In Proceedings of the 2019 IEEE/CVF International Conference on Computer Vision Workshop (ICCVW), Seoul, Republic of Korea, 27–28 October 2019; pp. 199–212. [Google Scholar] [CrossRef]
Model | Precision% | Recall% | [email protected]/% | [email protected]:0.95/% | Params/M |
---|---|---|---|---|---|
YOLOv8n | 88.5 | 80.2 | 90.1 | 52.5 | 3.2 |
YOLOv8n-CBAM | 89.2 | 81.3 | 91.4 | 54 | 3.9 |
YOLOv8-SE | 89.3 | 82.1 | 91.6 | 53.9 | 3.7 |
YOLOv8-MSSA * | 89.5 | 82 | 92.3 | 55 | 4.0 |
Model | Precision/% | Recall/% | [email protected]/% | [email protected]/% |
---|---|---|---|---|
YOLOv8n (base) | 87.5 | 80.2 | 89.4 | 53.5 |
SSD | 85 | 77.8 | 87.9 | 52 |
MASK_RCNN | 87 | 79 | 89.5 | 53.2 |
Faster-RCNN | 85.2 | 78.9 | 89.2 | 52.9 |
YOLOv8n-MSSA * (ours) | 89.5 | 82 | 90 | 55.9 |
Baseline | DyHead | GD | Prune | Precision/% | Recall/% | mAP @0.5/% | Params/M |
---|---|---|---|---|---|---|---|
✓ | × | × | × | 88.5 | 80.2 | 90.1 | 3.2 |
✓ | ✓ | × | × | 89.2 | 80.1 | 91.5 | 3.6 |
✓ | × | ✓ | × | 89 | 81.5 | 92.2 | 3.9 |
✓ | × | × | ✓ | 88.3 | 79.8 | 89.8 | 2.3 |
✓ | ✓ | ✓ | ✓ | 89.3 | 81.9 | 92 | 2.5 |
Disclaimer/Publisher’s Note: The statements, opinions and data contained in all publications are solely those of the individual author(s) and contributor(s) and not of MDPI and/or the editor(s). MDPI and/or the editor(s) disclaim responsibility for any injury to people or property resulting from any ideas, methods, instructions or products referred to in the content. |
© 2024 by the authors. Licensee MDPI, Basel, Switzerland. This article is an open access article distributed under the terms and conditions of the Creative Commons Attribution (CC BY) license (https://creativecommons.org/licenses/by/4.0/).
Share and Cite
Liu, P.; Zhang, Y.; Guo, G.; Ding, J. Enhanced Detection and Recognition of Road Objects in Infrared Imaging Using Multi-Scale Self-Attention. Sensors 2024, 24, 5404. https://doi.org/10.3390/s24165404
Liu P, Zhang Y, Guo G, Ding J. Enhanced Detection and Recognition of Road Objects in Infrared Imaging Using Multi-Scale Self-Attention. Sensors. 2024; 24(16):5404. https://doi.org/10.3390/s24165404
Chicago/Turabian StyleLiu, Poyi, Yunkang Zhang, Guanlun Guo, and Jiale Ding. 2024. "Enhanced Detection and Recognition of Road Objects in Infrared Imaging Using Multi-Scale Self-Attention" Sensors 24, no. 16: 5404. https://doi.org/10.3390/s24165404