Attention Network for Rail Surface Defect Detection via CASIoU-Guided Center-Point Estimation

Rail surface defect inspection based on machine vision faces challenges against the complex background with interference and severe data imbalance. To meet these challenges, we regard defect detection as a key-point estimation problem and present the attention neural network for rail surface defect detection via CASIoU-guided center-point estimation (CCEANN). CCEANN contains two crucial components. One is the stacked attention Hourglass backbone via cross-stage fusion of multi-scale features (CSFA-Hourglass), in which the convolutional block attention module with variable receptive fields (VRF-CBAM) is introduced, and a two-stage Hourglass structure balancing the network depth and feature fusion plays a key role. Furthermore, the CASIoU-guided center-point estimation head module (CASIoU-CEHM) integrating the delicate coordinate compensation mechanism regresses detection boxes flexibly to adapt to defects' large-scale variation, in which the proposed CASIoU loss, a loss regressing the consistency of Intersection-over-Union (IoU), central-point distance, area ratio, and scale ratio between the targeted defect and the predicted defect, achieves higher regression accuracy than state-of-the-art IoU-based losses. The experiments demonstrate that CCEANN outperforms competitive deep learning-based methods in four surface defect datasets.