# Rail-STrans: A Rail Surface Defect Segmentation Method Based on Improved Swin Transformer
## 2 Materials and Methods
### Local Perception Module (LPM)  
The camera-captured images contain valuable information about edges, contours, and textures. Therefore, this paper proposes using an improved Swin Transformer [47]-based feature extraction network to effectively extract rich contextual and long-term dependent information from the images. Figure 4 illustrates the structure of the improved Swin Transformer network. The network framework comprises a Patch Partition module and four cascaded layers that output four feature maps with varying resolutions in stages. The Patch Partition module is used to chunk the image, while the first layer contains a Linear Embedding layer, two Local Perception Modules (LPMs), and two Swin Transformer modules. The initial layer of the four cascaded layers consists of a Linear Embedding layer, two Local Perception Modules (LPMs), and two Swin Transformer modules. The Linear Embedding layer is utilized to modify the number of channels in the feature map. The subsequent layers consist of a Patch Merging layer, an even number of LPMs, and an even number of Swin Transformer modules. The Patch Merging layer downsamples the spatial resolution by 2-fold to reduce the feature resolution and adjust the number of channels. The Swin Transformer network framework has been improved by utilizing an even number of cascaded Swin Transformer modules. This allows for the division and merging  
Figure 4: Improved Swin Transformer network structure diagram. Figure 3: Dataset built in this paper. (**a**) Original picture. (**b**) Flip 45 degrees. (**c**) Flip horizontal. (**d**) Flip vertical. of the surface defect feature map, thus expanding the local receptive field to the global receptive field. The network framework consists of four levels, each with different-resolution output feature maps, S1, S2, S3, and S4, which are \(\frac{H}{4}\times\frac{H}{4}\), \(\frac{H}{8}\times\frac{H}{8}\), \(\frac{H}{16}\times\frac{H}{16}\), and \(\frac{H}{32}\times\frac{H}{32}\), respectively. These output feature maps have varying receptive fields, ranging from small to large, which aids in the segmentation of defect targets of different scales. The LPM framework, shown in Figure 5, obtains local context information of the image and reduces the number of network parameters. First, a set of feature vectors are reshaped into a spatial feature map using Reshape. Then, residual operations are performed on the outputs of DepthWise Separable Convolution (DWSC) and the SiLU activation module. Finally, the feature map is reduced to its original dimensions and output to the Swin Transformer module. Figure 6 shows the structure of the Swin Transformer module, which comprises Window Multi-head Self-Attention (W-MSA), Shift Window Multi-head Self-Attention (SW-MSA), and Multilayer Perceptron (MLP), each with Layer Normalization (LN). Figure 5: LPM framework. Figure 6: The architecture of the Swin Transformer block. LN is used to normalize the features in order to make the training process more stable; meanwhile, in order to reduce the computational complexity, the W-MSA module confines the computation of the self-attention to a single window, and the output of its \(l\)th level is  
\[z^{l}=\textit{W-MSA}\left(LN\left(z^{l-1}\right)\right)+z^{l-1} \tag{1}\]  
And the output of the MLP layer at level \(l\) is  
\[\hat{z}^{l}=MLP\left(LN\left(\hat{z}^{l}\right)\right)+\hat{z}^{l} \tag{2}\]  
In the next cascade module, SW-MSA adds a shift operation to W-MSA. This establishes information interaction between different windows without increasing computational effort. The output is as follows:  
\[\hat{z}^{l+1}=\textit{SW-MSA}\left(LN\left(z^{l}\right)\right)+z^{l} \tag{3}\]  
\[z^{l+1}=MLP\left(LN\left(\hat{z}^{l+1}\right)\right)+\hat{z}^{l+1} \tag{4}\]  
where \(\hat{z}^{l+1}\) represents the output of the SW-MSA module at level \(l+1\) and \(z^{l+1}\) represents the output of the MLP module at level \(l+1\). In accordance with the methodology outlined in Equation (5), the W-MSA is calculated in a manner analogous to that of the SW-MSA.