Cross-scale fusion and domain adversarial network for generalizable rail surface defect segmentation on unseen datasets

Surface quality control is a crucial part of rail manufacturing. Deep neural networks have shown impressive accuracy in rail surface defect segmentation under the assumption that the test images have the same distribution as the training images. However, in practice detection, the rail images exhibit variations in appearance and scale for different rail types and production conditions. Directly deploying the deep neural network on unseen images shows a performance degradation due to the distribution discrepancies of training images. To this end, we propose a cross-scale fusion and domain adversarial network (CFDANet) to improve the generalization ability of deep neural networks on unseen datasets. To alleviate the domain shift caused by defect scale differences, we design a dual-encoder to extract multi-scale features from images of different resolutions. Then, those features are adaptively fused through a cross-scale fusion module. For the domain shift caused by inconsistent rail appearance, we introduce transferable-aware domain adversarial learning to extract domain invariant features from different datasets. Moreover, we further propose a transferable curriculum to suppress the negative impact of images with low transferability. Experimental results show that our CFDANet can accurately segment defects in unseen datasets and surpass other state-of-the-art domain generalization methods in all five target domain settings. The source code is released at https://github.com/dotaball/railseg_dg .