Lightweight Algorithm for Rail Fastener Status Detection Based on YOLOv8n
Abstract
:1. Introduction
- Based on the characteristics of rail fasteners, an aspect ratio loss term is added to the EIOU loss function.
- Introducing ConvNeXt V2 modules, receptive field components and overlapped spatial attention (RCS-OSA) modules, and Efficient Rep networks to reconstruct the YOLOv8 backbone network, serving as the teacher model.
- Designing a lightweight C2f module to improve the YOLOv8 backbone network, serving as the student model.
- Improving the CWD feature knowledge distillation algorithm by designing an L2 loss function at the output end of the distillation to calculate the mean squared error loss between the student model and the teacher model outputs, enabling the student model to better learn the class features of the teacher model and enhance the model performance.
2. Methods
2.1. Regression Frame Loss Function Design
2.2. Teacher Model Design
2.2.1. ConvNeXt V2 Module
2.2.2. RCS-OSA Module
2.2.3. Efficient Rep Network
2.3. Student Model Design
2.4. Improved CWD Knowledge Distillation Algorithm
3. Experiment and Analysis
3.1. Experimental Data Set
3.2. Experimental Environment and Parameters
3.3. Evaluation Index
3.4. Comparative Experiment of Different Loss Functions
3.5. YOLOv8n-T Network Ablation Experiment
3.6. YOLOv8n-S Network Ablation Experiment
3.7. Knowledge Distillation Contrast Experiment
3.8. Comparison Experiment of Different Models
3.9. Comparison of Detection Effects before and after Improvement
4. Conclusions
- An EIOU+ localization loss function was designed based on the aspect ratio differences of railway fasteners, which accelerates the model’s convergence and enhances detection precision for more accurate results.
- The feature extraction network YOLOv8n-T was designed by integrating the ConvNeXt V2 module, RCS-OSA module, and Efficient Rep network to improve the original YOLOv8 backbone. This enhances the network’s feature extraction capability in complex environments, ensuring the better transmission of deep feature information to the Neck part and improving detection accuracy.
- A lightweight YOLOv8n-S model was designed, incorporating the C2f-Het module to improve the C2f module in the backbone network, achieving a lightweight design while maintaining the detection speed.
- An L2 loss function was designed to further optimize CWD knowledge distillation, resulting in the CWD+ knowledge distillation method. In this process, YOLOv8n-T was employed as the teacher model and YOLOv8n-S as the student model for model compression, leading to the training of a lightweight model.
Author Contributions
Funding
Data Availability Statement
Conflicts of Interest
References
- Wei, X.; Wei, D.; Suo, D.; Jia, L.; Li, Y. Multi-target defect identification for railway track line based on image processing and improved YOLOv3 model. IEEE Access 2020, 8, 61973–61988. [Google Scholar] [CrossRef]
- Zhuang, L.; Qi, H.Y.; Wang, T.G.; Zhang, Z. A Deep-learning powered Near-real-time Detection of Railway Track Major Components: A Two-stage Computer-vision-based Method. IEEE Internet Things J. 2022, 9, 18806–18816. [Google Scholar] [CrossRef]
- Ma, A.; Lv, Z.; Chen, X.; Li, L.; Qi, Y.; Zheng, S.; Chai, X. Pandrol track fastener defect detection based on local convolutional neural networks. Proc. Inst. Mech. Eng. Part I J. Syst. Control. Eng. 2021, 235, 1906–1915. [Google Scholar] [CrossRef]
- Prasongpongchai, T.; Chalidabhongse, T.H.; Leelhapantu, S. A vision-based method for the detection of missing rail fasteners. In Proceedings of the 2017 IEEE International Conference on Signal and Image Processing Applications (ICSIPA), Kuching, Malaysia, 12–14 September 2017; IEEE: Piscataway, NJ, USA, 2017; Volume 12–14, pp. 419–424. [Google Scholar]
- Wei, X.; Yang, Z.; Liu, Y.; Wei, D.; Jia, L.; Li, Y. Railway track fastener defect detection based on image processing and deep learning techniques: A comparative study. Eng. Appl. Artif. Intell. 2019, 80, 66–81. [Google Scholar] [CrossRef]
- Liu, J.; Teng, Y.; Shi, B.; Ni, X.; Xiao, W.; Wang, C.; Liu, H. A hierarchical learning approach for railway fastener detection using imbalanced samples. Measurement 2021, 186, 110240. [Google Scholar] [CrossRef]
- Redmon, J.; Divvala, S.; Girshick, R.; Farhadi, A. You Only Look Once: Unified, Real-Time Object Detection. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR), Las Vegas, NV, USA, 27–30 June 2016; pp. 779–788. [Google Scholar]
- Redmon, J.; Farhadi, A. YOLO9000: Better, Faster, Stronger. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR), Honolulu, HI, USA, 21–26 July 2017; pp. 6517–6525. [Google Scholar]
- Redmon, J.; Farhadi, A. Yolov3: An Incremental Improvement. arXiv 2018, arXiv:1804.02767. [Google Scholar]
- Chandran, P.; Asber, J.; Thiery, F.; Kumar, A.; Harsha, S.P. An Investigation of Railway Fastener Detection Using Image Processing and Augmented Deep Learning. Sustainability 2021, 13, 12051. [Google Scholar] [CrossRef]
- Qi, H.; Xu, T.; Wang, G.; Cheng, Y.; Chen, C. MYOLOv3-Tiny: A New Convolutional Neural Network Architecture for Real-Time Detection of Track Fasteners. Comput. Ind. 2020, 123, 103303. [Google Scholar] [CrossRef]
- Guo, F.; Qian, Y.; Shi, Y. Real-Time Railroad Track Components Inspection Based on the Improved YOLOv4 Framework. Autom. Constr. 2021, 125, 103596. [Google Scholar] [CrossRef]
- Wang, L.; Zang, Q.; Zhang, K.; Wu, L. A rail fastener defect detection algorithm based on improved YOLOv5. Proc. Inst. Mech. Eng. Part F J. Rail Rapid Transit. 2024, 238, 745–756. [Google Scholar] [CrossRef]
- Cai, Y.; He, M.; Tao, Q.; Xia, J.; Zhong, F.; Zhou, H. Fast Rail Fastener Screw Detection for Vision-Based Fastener Screw Maintenance Robot Using Deep Learning. Appl. Sci. 2024, 14, 3716. [Google Scholar] [CrossRef]
- Yang, Q.; Li, F.; Tian, H.; Li, H.; Xu, S.; Fei, J.; Wu, Z.; Feng, Q.; Lu, C. A new knowledge-distillation-based method for detecting conveyor belt defects. Appl. Sci. 2022, 12, 10051. [Google Scholar] [CrossRef]
- Lei, Y.; Chen, X.; Wang, Y.; Tang, R.; Zhang, B. A Lightweight Knowledge-Distillation-Based Model for the Detection and Classification of Impacted Mandibular Third Molars. Appl. Sci. 2023, 13, 9970. [Google Scholar] [CrossRef]
- Zhou, P.; Aysa, A.; Ubul, K. Research on knowledge distillation algorithm based on Yolov5 attention mechanism. Expert Syst. Appl. 2024, 240, 122553. [Google Scholar]
- Zheng, Z.; Wang, P.; Liu, W.; Li, J.; Ye, R.; Ren, D. Faster and Better Learning for Bounding Box Regression. arXiv 2020, arXiv:1911.08287. [Google Scholar] [CrossRef]
- Zhang, Y.F.; Ren, W.; Zhang, Z.; Jia, Z.; Wang, L.; Tan, T. Focal and efficient IOU loss for accurate bounding box regression. Neurocomputing 2022, 506, 146–157. [Google Scholar] [CrossRef]
- Yu, J.; Jiang, Y.; Wang, Z.; Cao, Z.; Huang, T. Unitbox: An advanced object detection network. In Proceedings of the 24th ACM International Conference on Multimedia, Amsterdam, The Netherlands, 15–19 October 2016; pp. 516–520. [Google Scholar]
- Woo, S.; Debnath, S.; Hu, R.; Chen, X.; Liu, Z.; Kweon, I.S.; Xie, S. Convnext v2: Co-designing and scaling convnets with masked autoencoders. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, Vancouver, BC, Canada, 17–24 June 2023; pp. 16133–16142. [Google Scholar]
- Zhang, X.; Zhou, X.; Lin, M.; Sun, J.; Inc, M. Shufflenet: An extremely efficient convolutional neural network for mobile devices. In Proceedings of the IEEE conference on computer vision and pattern recognition, Salt Lake City, UT, USA, 18–23 June 2018; pp. 6848–6856. [Google Scholar]
- Kang, M.; Ting, C.M.; Ting, F.F.; Greenspan, H.; Madabhushi, A.; Mousavi, P.; Salcudean, S.; Duncan, J.; Syeda-Mahmood, T.; Taylor, R. RCS-YOLO: A fast and high-accuracy object detector for brain tumor detection. In Proceedings of the International Conference on Medical Image Computing and Computer-Assisted Intervention, Vancouver, BC, Canada, 8–12 October 2023; Springer Nature: Cham, Switzerland, 2023; pp. 600–610. [Google Scholar]
- Weng, K.; Chu, X.; Xu, X.; Huang, J.; Wei, X. Efficientrep: An efficient Repvgg-style convnets with hardware-aware neural network design. arXiv 2023, arXiv:2302.00386. [Google Scholar]
- Ding, X.; Zhang, X.; Ma, N.; Han, J.; Ding, G.; Sun, J. Repvgg: Making vgg-style convnets great again. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, Nashville, TN, USA, 20–25 June 2021; pp. 13733–13742. [Google Scholar]
- Singh, P.; Verma, V.K.; Rai, P.; Namboodiri, V.P. Hetconv: Heterogeneous kernel-based convolutions for deep cnns. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, Long Beach, CA, USA, 15–20 June 2019; pp. 4835–4844. [Google Scholar]
- Shu, C.; Liu, Y.; Gao, J.; Yan, Z.; Shen, C. Channel-wise knowledge distillation for dense prediction. In Proceedings of the IEEE/CVF International Conference on Computer Vision, Virtual, 11–17 October 2021; pp. 5311–5320. [Google Scholar]
- Li, T.; Wang, J.; Yao, K. Visibility enhancement of underwater images based on active polarized illumination and average filtering technology. Alex. Eng. J. 2022, 61, 701–708. [Google Scholar] [CrossRef]
- Rizk-Allah, R.M.; Hassanien, A.E. A comprehensive survey on the sine–cosine optimization algorithm. Artif. Intell. Rev. 2023, 56, 4801–4858. [Google Scholar] [CrossRef]
- Padilla, R.; Netto, S.L.; Da Silva, E.A. A Survey on Performance Metrics for Object-Detection Algorithms. In Proceedings of the 2020 International Conference on Systems, Signals and Image Processing (IWSSIP), Niteroi, Brazil, 1–3 July 2020; pp. 237–242. [Google Scholar]
- Shah, T. Measuring Object Detection Models—MAP—What Is Mean Average Precision. Tarang. Shah.-Blog. 2018, 26, 104332. [Google Scholar]
- Tong, Z.; Chen, Y.; Xu, Z.; Yu, R. Wise-IoU: Bounding box regression loss with dynamic focusing mechanism. arXiv 2023, arXiv:2301.10051. [Google Scholar]
- Siliang, M.; Yong, X. MPDIoU: A loss for efficient and accurate bounding box regression. arXiv 2023, arXiv:2307.07662. [Google Scholar]
Loss Function | P% | R% | mAP@50/% | mAP@50–95/% |
---|---|---|---|---|
CIOU | 88.7 | 86.6 | 93.6 | 66.5 |
EIOU | 91.2 | 88.2 | 93.3 | 67.3 |
WIOU | 90.5 | 88.4 | 93.8 | 68.8 |
MPDIOU | 91.2 | 90 | 94.2 | 68.6 |
EIOU+ | 88.6 | 92.3 | 94.5 | 69.1 |
YOLOv8n | P/% | R/% | mAP@50/% | mAP@ 50–90/% | GFLOPS | FPS | |||
---|---|---|---|---|---|---|---|---|---|
Conv NeXt V2 | RCS-OSA | Efficient Rep | EIoU+ | ||||||
88.7 | 88.6 | 93.6 | 66.5 | 8.4 | 36.5 | ||||
√ | 87.5 | 92.5 | 94.8 | 68.4 | 8.3 | 35.6 | |||
√ | 91.2 | 89.3 | 94.9 | 68.4 | 11.1 | 37.8 | |||
√ | 91.9 | 88.9 | 93.7 | 68.3 | 8.5 | 43.7 | |||
√ | √ | 89.8 | 92 | 94.9 | 68.7 | 11 | 31.9 | ||
√ | √ | 90.9 | 90.5 | 94.5 | 68.6 | 8.5 | 35.9 | ||
√ | √ | 91.4 | 90.5 | 94.6 | 67.7 | 11.2 | 37.7 | ||
√ | √ | √ | 89.8 | 90.8 | 94.2 | 67.7 | 11.2 | 31.8 | |
√ | 88.6 | 92.3 | 94.5 | 69.1 | 8.4 | 43.9 | |||
√ | √ | √ | √ | 90.9 | 90.7 | 95.6 | 69.6 | 11.2 | 31.8 |
YOLOv8n | P/% | R/% | mAP@ 50/% | mAP@ 50–90/% | GFLOPS | FPS | |
---|---|---|---|---|---|---|---|
C2f-Het | EIoU+ | ||||||
88.7 | 88.6 | 93.6 | 66.5 | 8.4 | 36.6 | ||
√ | 92.5 | 86.6 | 94.6 | 67.9 | 7.6 | 43.5 | |
√ | 88.6 | 92.3 | 94.5 | 69.1 | 8.4 | 43.9 | |
√ | √ | 90.6 | 89.5 | 95.2 | 68.4 | 7.6 | 41.6 |
Temperature | P/% | R/% | mAP@50/% | mAP@50–90/% |
---|---|---|---|---|
= 1.0 | 91.8 | 91.1 | 96.3 | 69.4 |
= 2.0 | 91.5 | 90.6 | 95.6 | 68.9 |
= 3.0 | 91.5 | 90.6 | 95.6 | 68.9 |
YOLOv8n-T+ YOLOv8n-S | P/% | R/% | mAP@ 50/% | mAP@ 50–90/% | GFLOPS | FPS | |
---|---|---|---|---|---|---|---|
CWD | CWD+ | ||||||
√ | 89.5 | 90.4 | 95.3 | 68.3 | 7.3 | 42.5 | |
√ | 91.8 | 91.1 | 96.3 | 69.4 | 7.3 | 42.7 |
Model | P/% | R/% | mAP@ 50/% | mAP@ 50–90/% | GFLOPS | FPS |
---|---|---|---|---|---|---|
YOLOv3-tiny | 88.4 | 84.8 | 91.4 | 60.8 | 13 | 42.5 |
YOLOv5s | 92.9 | 86.7 | 93.4 | 64.4 | 16 | 27.3 |
YOLOv8n | 88.7 | 88.6 | 93.6 | 66.5 | 8.4 | 36.6 |
YOLOv8n-T | 90.9 | 90.7 | 95.6 | 69.6 | 11.2 | 31.8 |
YOLOv8n-S | 90.6 | 89.5 | 95.2 | 68.4 | 7.6 | 41.6 |
YOLOv8n-CWD | 89.5 | 90.4 | 95.3 | 68.3 | 7.3 | 42.5 |
YOLOV8n-CWD+ | 91.8 | 91.1 | 96.3 | 69.4 | 7.3 | 42.7 |
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
Zhang, X.; Shen, B.; Li, J.; Ruan, J. Lightweight Algorithm for Rail Fastener Status Detection Based on YOLOv8n. Electronics 2024, 13, 3399. https://doi.org/10.3390/electronics13173399
Zhang X, Shen B, Li J, Ruan J. Lightweight Algorithm for Rail Fastener Status Detection Based on YOLOv8n. Electronics. 2024; 13(17):3399. https://doi.org/10.3390/electronics13173399
Chicago/Turabian StyleZhang, Xingsheng, Benlan Shen, Jincheng Li, and Jiuhong Ruan. 2024. "Lightweight Algorithm for Rail Fastener Status Detection Based on YOLOv8n" Electronics 13, no. 17: 3399. https://doi.org/10.3390/electronics13173399