CBAM-SwinT-BL: Small Rail Surface Defect Detection Method Based on Swin Transformer with Block Level CBAM Enhancement

Jiayi Zhao,Alison Wun-lam Yeung,Ali Muhammad,Songjiang Lai,Vincent To-Yee NG
2024-10-02
Abstract:Under high-intensity rail operations, rail tracks endure considerable stresses resulting in various defects such as corrugation and spellings. Failure to effectively detect defects and provide maintenance in time would compromise service reliability and public safety. While advanced models have been developed in recent years, efficiently identifying small-scale rail defects has not yet been studied, especially for categories such as Dirt or Squat on rail surface. To address this challenge, this study utilizes Swin Transformer (SwinT) as baseline and incorporates the Convolutional Block Attention Module (CBAM) for enhancement. Our proposed method integrates CBAM successively within the swin transformer blocks, resulting in significant performance improvement in rail defect detection, particularly for categories with small instance sizes. The proposed framework is named CBAM-Enhanced Swin Transformer in Block Level (CBAM-SwinT-BL). Experiment and ablation study have proven the effectiveness of the framework. The proposed framework has a notable improvement in the accuracy of small size defects, such as dirt and dent categories in RIII dataset, with mAP-50 increasing by +23.0% and +38.3% respectively, and the squat category in MUET dataset also reaches +13.2% higher than the original model. Compares to the original SwinT, CBAM-SwinT-BL increase overall precision around +5% in the MUET dataset and +7% in the RIII dataset, reaching 69.1% and 88.1% respectively. Meanwhile, the additional module CBAM merely extend the model training speed by an average of +0.04s/iteration, which is acceptable compared to the significant improvement in system performance.
Computer Vision and Pattern Recognition
What problem does this paper attempt to address?
### What problem does this paper attempt to solve? This paper aims to address the issue of detecting small defects on the surface of railway tracks. In high-intensity railway operations, tracks endure significant stress, leading to various defects such as corrugation and spalling. If these defects are not detected and maintained in a timely and effective manner, it will affect service reliability and public safety. Although some advanced models have been developed in recent years, efficiently identifying small-scale railway defects (especially surface defects like dirt or dents) has not been fully studied. To this end, this research proposes a method based on Swin Transformer combined with Convolutional Block Attention Module (CBAM) enhancement, named CBAM-SwinT-BL (CBAM-Enhanced Swin Transformer in Block Level). Experiments and ablation studies show that this framework significantly improves performance in railway defect detection, particularly in small instance categories. For example, in the RIII dataset, the mAP-50 for dirt and dents categories increased by 23.0% and 38.3%, respectively, and in the MUET dataset, the mAP-50 for the dents category increased by 13.2%. Compared to the original Swin Transformer, CBAM-SwinT-BL improved overall accuracy by approximately 5% on the MUET dataset and 7% on the RIII dataset, reaching 69.1% and 88.1%, respectively. Additionally, the CBAM module only increases the model training speed by 0.04 seconds per iteration, which is an acceptable cost relative to the significant improvement in system performance.