The purpose of sample expansion is to provide an adequate number of samples for the learning and training of the deep network. In the basic generative adversarial network model, the generator receives a random sample space with a Gaussian distribution and noise data, while the discriminator receives the samples generated by the generator and real data samples. Both of them use the loss function to perform backpropagation of gradients to update the parameters of their own network models. However, this model uses random samples with a Gaussian distribution and noise data, ignoring the characteristics of the information distribution of the actual application objects, resulting in poor quality of the generated samples. In view of this, during the sample expansion stage of this paper, noise data is added based on the AdaIN mechanism, enabling the generator to generate rail surface condition images according to the image size. Moreover, independent Gaussian noise is input into each layer, which only causes slight changes in the visual features.

Among them, AdaIN is defined as

$\text{AdaIN}\left(x_i,y\right)=\delta \left(y\right)\left(\frac{x_i-\mu \left(x_i\right)}{\sigma \left(x_i\right)}\right)+\mu \left(y\right)$ (1)

In the formula, μ(x_i) and σ(x_i) represent the mean and variance of the real samples respectively, μ(y) and δ(y) represent the mean and variance of the latent sample space respectively.

The data network generated by using formula (1) can better fit the data distribution of the rail surface images and generate image information that is more in line with the real samples. For the specific network structure, you can refer to reference [12] .

The features of the rail surface condition are subtle, randomly located, and greatly interfered by background noise, which leads to the small differences in the feature subsets extracted by the traditional convolutional neural network and the unobvious weight comparison. Based on this, this paper designs a feature extraction module embedded with a multi-dimensional attention mechanism. The specific approach is as follows: channel attention is added in the channel dimension to highlight the information of the feature channels; attention is applied in the spatial dimension to strengthen the information in the spatial dimension. The module is shown in Figure 3 .

As shown in Figure 3 , in the process of enhancing the rail surface condition features, the input consists of rail surface condition image samples. These samples are first processed by a convolutional neural network to obtain feature vectors containing high-dimensional semantic information. Then, the weighted attention mechanism parameters are calculated through the channel attention mechanism and the spatial attention mechanism. Subsequently, the feature maps composed of feature vectors are multiplied element-by-element with the attention weights obtained from the attention mechanism, and finally, the output result after attention weighting is obtained. Among them, the expressions of the channel attention mechanism and the spatial attention mechanism are as follows

$Mc\left(x\right)=d\left(MLP\left(\text{Conv}5\left(x\right)\right)+MLP\left(\text{Conv}5\left(x\right)\right)\right)$ (2)

$Ms(x)=d\left(\text{Conv}\left(Mc\left(x\right)\right)\oplus Mc\left(x\right)\right)$ (3)

$At\left(x\right)=\text{Conv}5\left(x\right)\otimes Mc\left(x\right)$ (4)

Among them, Conv5 is the feature space formed after passing through five convolutional layers. Mc and Ms represent the channel attention weight and the spatial attention weight respectively in sequence, and At is the weighted feature space. By using the channel attention mechanism, the spatial dimension can be compressed, the features of the rail surface condition can be prominently highlighted, and high weights can be assigned to them. By applying the spatial attention mechanism, the channel dimension can be compressed, and the spatial position information can be supplemented to the model, so as to improve the utilization rate of the feature vectors.

The fully connected layer is generally used in the decision-making layer of a neural network. This module is characterized by a large number of parameters and difficult training. In this paper, the idea of residuals is introduced into the spinal fully connected layer, and a residual spinal fully connected layer is proposed. This module innovatively improves the input-output connection and adds a residual channel to alleviate the overfitting problem of the spinal fully connected layer. Moreover, the activation function is reset based on the quantity characteristics of the data set, making this module more suitable for the small sample data set of rail surface condition images. The model structure is shown in Figure 4 .

Among them, the feature A_t extracted by the module with the attention mechanism introduced is split into two parts, A1 and A2.

$A1=A2=\left\{A_{t}1,A_{t}2,\cdots ,A_{t}640\right\}$ (5)

The spinal neural network in this paper is composed of four modules, namely Sp1, Sp2, Sp3, and Sp4. Each module contains a dropout layer, a linear layer, and an activation function. The settings of the dropout layer and the activation function are the same, and the output of the linear layer is 512 for all of them. The difference lies in that the input size of the linear layer in the Sp1 module is 1024, while the input size of the linear layer in the Sp2, Sp3, and Sp4 modules is 1536. Such a setting is to keep the feature dimensions input to the fully connected layer consistent

$Sp1={\sigma }_2\left(\text{Linear}\left(\text{Droop}\left(A_1\right)\right)\right)$ (6)

$Sp2=\sigma \left(\text{Linear}\left(\text{Droop}\left(Sp1\oplus A2\right)\right)\right)$ (7)

$Sp3=\sigma \left(\text{Linear}\left(\text{Droop}\left(Sp2\oplus A1\right)\right)\right)$ (8)

$Sp4=\sigma \left(\text{Linear}\left(\text{Droop}\left(Sp3\oplus A2\right)\right)\right)$ (9)

$FC=\text{linear}\left(Sp1+Sp2+Sp3+Sp4\right)+\text{Maxpool}\left(At\left(X\right)\right)$ (10)

Here, Droop represents a random zeroing operation, Linear stands for the linear layer, and σ is the activation function of the spinal neural network. The output of each spinal block, when combined with the feature vector after max pooling, yields a fine-tuned feature vector of size 1 × 1 × 2048. Eventually, this enables the accomplishment of the pattern recognition task for rail surface condition images.

The dataset of rail surface condition images used in this paper was captured by a SONY-ILCE-6000 camera at different time periods in a local train depot. A total of 210 rail surface condition images were selected and the image categories were labeled. Among them, there were 70 images each for dry rail surfaces, wet rail surfaces, and oil-contaminated rail surfaces. Since the original images captured by the camera contain a large amount of background information and noise interference, which increases the difficulty of model recognition, pre-processing operations were carried out on the rail surface condition dataset. These operations include median denoising, geometric correction, and extraction of rail surface area images, so as to improve the robustness and recognition accuracy of the network.

To objectively measure the performance of the generative adversarial network and address the problem of insufficient training samples, this paper constructs a rail surface condition image dataset, which covers various forms of data enhancement: the original dataset without enhancement (Non), the dataset processed by geometric transformation (GT), the dataset processed by pixel transformation (PT), and the dataset enhanced by the generative adversarial network (GAN). Meanwhile, for the convenience of network training and testing, all images were uniformly reconstructed to a specification of 256 × 256 pixels. The specific distribution of the enhanced dataset is shown in Table 1 . Finally, the self-built dataset was approximately divided into a training set and a test set at a ratio of 4:1.

The experimental program in this paper uses PyTorch as the operating platform, and the running environment is configured as follows: The processor is Intel Core i9-9900K, the running memory is 20G, and the graphics card is NVIDIA RTX 4090. The code operating environment is Torch = 1.90, Python = 3.12, and the CUDA version is CUDA12.4. Stochastic Gradient Descent (SGD) is selected as the Optimizer, the batch size is set to 32, and the training is carried out for 50 iterations.

The evaluation indicators of the experiments in this paper include accuracy, recall rate, precision rate, and F1 score. Among them, the accuracy refers to the proportion of correctly predicted samples in the total number of samples; the recall rate R refers to the proportion of positive samples that are correctly predicted; the precision rate P refers to the proportion of truly positive samples among the samples predicted as positive; the F1 score comprehensively considers the recall rate and precision rate, and it is the weighted harmonic mean of the two.

In the above three datasets, Resnet is used as the feature extraction module to verify whether the data augmentation method is effective. The results are shown in Table 2 .

It can be seen from Table 2 that data augmentation can expand the number of training samples, improve sample diversity, and expand the spatial distribution of samples. Therefore, in all models, whether using geometric transformation, pixel transformation, or generative adversarial networks to expand the dataset, the model performance can be improved and the error rate can be reduced. Comparing geometric transformation and pixel transformation, the rail surface adhesion media are mostly distributed in blocks and their colors are darker than the background, which is similar to the noise added by pixel transformation and can easily mislead the model into making classification errors. So, the dataset expansion effect of geometric transformation is slightly better than that of pixel transformation. When applying the generative adversarial network, a realistic and class-balanced dataset can be generated based on real operation and maintenance data. The recognition accuracy on the ResNet pre-trained network reaches 91.35%, which is better than that of the datasets using geometric transformation and pixel transformation. Compared with traditional data augmentation methods, the generative adversarial network can generate almost an infinite amount of data through random seeds, meeting the large-scale data requirements of deep convolutional neural networks, alleviating the model’s data dependence, and enhancing the anti-overfitting ability.

To verify the impact of the improved fully-connected layer on the model performance, the fully-connected layer in Resnet-50 is replaced with the improved structure proposed in this paper. The performance differences between the spinal fully-connected layer and the residual spinal fully-connected layer are tested, and a comparative experiment is conducted on these two models to analyze the performance differences of the model before and after replacement. In the experiment, except for the fully-connected layer being different, the rest of the settings are the same, and the GAN dataset images are used for verification. The results are shown in Table 3 .

In the experimental settings, “Non” represents Resnet50 using a fully-connected layer, and “Spinal” refers to the situation where the fully-connected layer of Resnet50 is replaced with a spinal fully-connected layer. Although the spinal neural network can optimize the performance of some network models to a certain extent, when applied to Resnet50, it will lead to an increase in the number of model parameters. “Spinal-res” is the model form after replacing the fully-connected layer in Resnet50 with a residual spinal fully-connected layer. By looking at the data in Table 3 , it can be found that compared with the original fully-connected layer of Resnet-50, the residual spinal fully-connected layer improves the model’s recognition accuracy by 1.32%, precision by 1.12%, recall by 1.45%, and specificity by 1.26%. This means that the residual spinal fully-connected layer has better recognition ability for the rail surface state image features. The reason is that its dual-branch decision-making network has both local and global perspectives, which can effectively reduce the interference of local information, enhance the model’s stability and anti-interference ability, and improve the model’s robustness.

The dimensions and weights of the feature layers can affect the network training results. During the training of the Resnet network, since the influence of feature information in different dimensions on the model recognition results varies, an attention mechanism module is added after feature extraction. In this experiment, two attention mechanism modules, the single-channel attention mechanism (SE) and the multi-channel attention mechanism (CBAM), are set up, and the GAN dataset images are used for verification. The results are shown in Table 4 .

It can be seen from Table 4 that in terms of the attention mechanism, when it is set to SE, compared with the situation without the attention mechanism module, the model’s accuracy rate is increased by 1.54%, the precision rate is increased by 1.52%, the recall rate is increased by 1.79%, and the specificity rate is increased by 1.6%. When the attention mechanism is set to CBAM, compared with the case without this module, the accuracy rate is increased by 2.11%, the precision rate is increased by 2.36%, the recall rate is increased by 2.71%, and the specificity rate is increased by 2.99%. Obviously, the multi-channel attention mechanism (CBAM) has better performance. It can assign corresponding weights according to the importance of different feature information, thereby strengthening the expression of the rail surface state features and improving the overall performance of the model.