A steel rail small sample defect identification method, system, device and medium

Abstract

The application relates to the field of defect identification, and discloses a rail small sample defect identification method, system, device and medium, the method comprising: constructing a feature extraction network, which comprises a convolutional neural network backbone and an adaptive feature weighting module, the module recalibrates the initial features based on a channel and space double attention mechanism, and outputs a weighted feature vector; prototype feature vectors of each category are calculated based on the weighted feature vectors of the support set samples; the distances between the weighted feature vectors of the query samples and the prototype vectors of each category are calculated, a probability distribution is obtained through softmax, and a meta-training loss is calculated based on the difference with the real label; the network parameters are optimized through gradient descent, and the feature extraction network with small sample learning ability is obtained through iterative training; in the meta-test stage, the support set and the query set containing new defect categories are input into the trained network to obtain a classification result. The application can accurately identify rail surface defects under very few labeled samples.

Description

Technical Field

[0001] This application relates to the field of defect identification, and in particular to a method, system, equipment and medium for identifying small samples of defects in rails. Background Technology

[0002] Traditional rail defect detection mainly relies on manual inspection or large-scale specialized flaw detection equipment, which suffers from low efficiency, high cost, and difficulty in achieving comprehensive coverage. In recent years, deep learning-based methods have become mainstream, but they depend on large-scale labeled data. In rail defect scenarios, high-risk defect samples are scarce, labeling costs are high, and there is a risk of long-tailed distribution and the emergence of new defects, making it difficult to construct sufficient training data.

[0003] To address the problem of scarce samples, existing technologies include data augmentation, transfer learning, and few-shot learning methods. Among these, few-shot learning methods, such as prototype networks, show potential by classifying prototype features from a small number of samples. However, rail images have complex backgrounds and subtle defect areas, making it difficult for standard prototype networks to effectively distinguish defect features from background interference when calculating prototypes. This results in insufficient discriminative power of the prototype representation, severely limiting classification accuracy when samples are extremely limited.

[0004] Therefore, existing technologies lack an intelligent recognition method that can adaptively focus on key defect features and suppress background interference with a very small number of samples, thereby constructing a prototype with high discriminative power. Summary of the Invention

[0005] To address the problems of insufficient discriminative power of prototype features, susceptibility to background interference, and weak recognition generalization ability of existing rail defect detection methods when samples are scarce, this application provides a method, system, equipment, and medium for identifying rail defects in small samples.

[0006] This application discloses a method for identifying small-sample defects in steel rails, which includes: Step 1: Construct a meta-training task set, which consists of multiple few-shot recognition tasks. Each few-shot recognition task is constructed as follows: From the meta-training set containing images of multiple types of rail defects, randomly select N different defect categories; for each of the N selected defect categories, randomly select K labeled samples from each defect category to form a support set; simultaneously, from the same N defect categories, randomly select Q labeled samples for each category, and the selected samples are not repeated with the samples in the support set to form a query set. Step 2: Construct a feature extraction network, which includes a convolutional neural network (CNN) backbone and an adaptive feature weighting module connected to the CNN backbone. The CNN backbone is used to extract the initial feature representation of the rail image input to it. The adaptive feature weighting module is used to recalibrate the initial feature representation based on a dual channel and spatial attention mechanism to enhance key features related to defects and suppress irrelevant background features, outputting a weighted feature vector. The key features include crack features, spalling features, or indentation features on the rail surface. The initial feature representation refers to the multi-channel feature map of the rail image extracted by the CNN backbone, used to characterize the original visual features related to defects in the rail image. Step 3: For each of the few sample recognition tasks, input the samples in the support set into the feature extraction network to obtain the weighted feature vector of each sample; based on the weighted feature vector of all samples of the same category, calculate the prototype feature vector of that category in the feature space; Step 4: Input the samples in the query set into the feature extraction network to obtain the weighted feature vector of each query sample; calculate the distance between the weighted feature vector of each query sample and the prototype feature vector of all categories obtained in Step 3. Step 5: Based on the distance, calculate the probability distribution of each query sample belonging to each category using the softmax function, and calculate the meta-training loss based on the difference between the probability distribution and the true label of the query sample; Step 6: Optimize the parameters of the feature extraction network using the gradient descent algorithm, minimize the meta-training loss, and iterate through multiple small sample recognition tasks to enable the feature extraction network to learn the ability to recognize new defect categories from a small number of samples, thus obtaining a trained feature extraction network. Step 7: In the meta-testing phase, the support set and query set containing the new defect categories are input into the trained feature extraction network to obtain the sample classification results in the query set containing the new defect categories.

[0007] Furthermore, the adaptive feature weighting module is a channel-space dual attention module, which includes a channel attention submodule and a spatial attention submodule connected in sequence; The channel attention submodule is used to generate a weight vector based on the initial feature representation, and then multiply the weight vector with the initial feature representation channel by channel to generate a channel-weighted feature representation. The spatial attention submodule is used to generate a spatial dimension weight map based on the channel-weighted feature representation, and to multiply the channel-weighted feature representation position by position using the spatial dimension weight map to output the weighted feature vector.

[0008] Further, the generation of the weight vector includes: The initial feature representation is subjected to global average pooling and global max pooling respectively to obtain two channel descriptors; The two channel descriptors are input into a shared multilayer perceptron network, which includes a dimensionality reduction layer and a recovery layer. The two features output by the multilayer perceptron network are added element by element, and then passed through the sigmoid activation function to generate the weight vector.

[0009] Furthermore, the generation of the spatial dimension weight graph includes: The channel-weighted feature representations are subjected to average pooling and max pooling along the channel dimension to obtain two spatial feature maps. The two spatial feature maps are concatenated along the channel dimension; The concatenated feature maps are then subjected to convolution operations and then activated by the Sigmoid function to generate a weight map with spatial dimensions.

[0010] Furthermore, in step 3, the formula for calculating the prototype feature vector of the nth type of defect in the feature space is:

[0011] in, Let n be the prototype feature vector of the nth type of defect in the feature space. Indicates the class n in the support set. One labeled sample, express The weighted feature vector obtained after passing through the feature extraction network has a value of n ranging from 1 to N, and K is the total number of samples of the nth type of defect.

[0012] Furthermore, in step 5, the calculated distance is either the Euclidean distance or the inverse of the cosine similarity; when using the Euclidean distance, the probability that the query sample x belongs to the nth type of defect is calculated. The calculation formula, expressed using the softmax function, is as follows:

[0013] Where y is the class label variable of the sample; Given model parameters θ, the probability that the category label y of the query sample x is the nth type of defect; The weighted feature vector obtained after the query sample x has passed through the feature extraction network is... Let n be the prototype feature vector of the nth type of defect in the feature space. express and The Euclidean distance between them; The summation index variable has a range of 1 to N, and exp is an exponential function.

[0014] Furthermore, the meta-training loss function in step 5 is the negative log-likelihood loss, and its calculation formula is as follows:

[0015] Where T is the total number of samples in the query set. This represents the t-th sample in the query set; for The true label, For query samples Belongs to real labels The probability of the category, This is the training loss.

[0016] This application also discloses a rail small-sample defect identification system for implementing the above-described method, comprising: The task construction module is used to construct a meta-training task set, which consists of multiple few-sample recognition tasks. These tasks are generated iteratively from a meta-training set containing images of multiple types of rail defects. Each few-sample recognition task includes a support set and a query set. The support set contains N types of defect samples, each type of defect sample having K labeled samples. The query set contains labeled samples that are not repeated from the labeled samples in the support set, and each sample in the query set belongs to one of the N types of defects in the support set. A network construction module is used to construct a feature extraction network, which includes a convolutional neural network (CNN) backbone and an adaptive feature weighting module connected to the CNN backbone. The CNN backbone is used to extract an initial feature representation of the rail image input to it. The adaptive feature weighting module is used to recalibrate the initial feature representation based on a dual channel and spatial attention mechanism to enhance key features related to defects and suppress irrelevant background features, outputting a weighted feature vector. The key features include crack features, spalling features, or indentation features on the rail surface. The initial feature representation refers to the multi-channel feature map of the rail image extracted by the CNN backbone, used to characterize the original visual features related to defects in the rail image. The feature vector calculation module is used to input the samples in the support set into the feature extraction network for each of the few sample recognition tasks to obtain the weighted feature vector of each sample; and to calculate the prototype feature vector of the category in the feature space based on the weighted feature vector of all samples of the same category. The distance calculation module is used to input the samples in the query set into the feature extraction network to obtain the weighted feature vector of each query sample; and to calculate the distance between the weighted feature vector of each query sample and the prototype feature vector of all categories. The meta-training loss calculation module is used to calculate the probability distribution of each query sample belonging to each category based on the distance using the softmax function, and to calculate the meta-training loss based on the difference between the probability distribution and the true label of the query sample. The model training module is used to optimize the parameters of the feature extraction network through the gradient descent algorithm, minimize the meta-training loss, and through multiple iterations of the few sample recognition tasks, enable the feature extraction network to learn the ability to recognize new defect categories from a small number of samples, thus obtaining a trained feature extraction network. The classification and recognition module is used in the meta-testing stage to input the support set and query set containing the new defect category into the trained feature extraction network to obtain the sample classification results in the query set containing the new defect category.

[0017] This application also discloses a computer-readable storage medium having a computer program stored thereon, which, when executed by a processor, implements the method described above.

[0018] This application also discloses an electronic device, including a memory, a processor, and a computer program stored in the memory and executable on the processor, wherein the processor executes the computer program to implement the method described above.

[0019] Compared with the prior art, the technical solution provided in this application has the following significant advantages: 1. It fundamentally solves the challenges of model training and generalization under small sample conditions. Through the meta-learning paradigm, the model acquires the ability to learn new categories from a small number of samples during the training phase. When faced with entirely new defect categories, there is no need for retraining; high-precision identification can be achieved quickly simply by calculating the prototype of the support set samples, breaking through the inherent constraints of data scarcity on deep learning applications.

[0020] 2. Significantly improves the prototype discriminative power and model robustness of feature representations under small sample conditions. By introducing an adaptive feature weighting module (such as a dual channel and spatial attention mechanism), the model can automatically focus on subtle defect areas in rail images and suppress interference from complex backgrounds such as ballast, fasteners, and uneven lighting. This makes the prototype feature vectors calculated from a small number of samples more "pure" and have higher discriminative power.

[0021] 3. It endows the system with powerful rapid adaptation and generalization capabilities. The model learns general "how to learn" capabilities, rather than a fixed classifier. Therefore, the system is not only effective for defect categories seen during training, but can also quickly adapt to and identify new defect types that have not appeared at all during the training phase. This meets the urgent need for real-time early warning of unknown defects in on-site inspections, and improves the system's practicality and foresight.

[0022] 4. A feasible and efficient training and recognition framework is provided. By constructing a meta-learning task that simulates real-world small-sample scenarios and integrating an adaptive feature enhancement mechanism, a complete "learning-adaptation-recognition" technical solution is formed. This solution can achieve high-precision and robust rail defect recognition with minimal labeled samples. Attached Figure Description

[0023] To more clearly illustrate the technical solutions in the embodiments of this application, the accompanying drawings used in the description of the embodiments will be briefly introduced below. Obviously, the accompanying drawings described below are only some embodiments recorded in the embodiments of this application. For those skilled in the art, other drawings can be obtained based on these drawings.

[0024] Figure 1 This is a flowchart illustrating a method for identifying small-sample defects in rails according to an embodiment of this application. Figure 2 This is a schematic diagram of a channel-space dual attention module structure according to an embodiment of this application; Figure 3 This is a schematic diagram of the meta-learning task sampling in an embodiment of this application. Detailed Implementation

[0025] The present application will be further described in conjunction with the accompanying drawings and embodiments. The described embodiments are only some, not all, of the embodiments of the present application. All other embodiments obtained by those skilled in the art should fall within the protection scope of the embodiments of the present application.

[0026] See Figure 1 This application provides an embodiment of a method for identifying small-sample defects in rails, comprising: Step 1: Construct a meta-training task set, which consists of multiple few-shot recognition tasks. Each few-shot recognition task is constructed as follows: From the meta-training set containing images of multiple types of rail defects, randomly select N different defect categories; for each of the N selected defect categories, randomly select K labeled samples from each defect category to form a support set; simultaneously, from the same N defect categories, randomly select Q labeled samples for each category, and the selected samples are not repeated with the samples in the support set to form a query set. Step 2: Construct a feature extraction network, which includes a convolutional neural network (CNN) backbone and an adaptive feature weighting module connected to the CNN backbone. The CNN backbone is used to extract the initial feature representation of the rail image input to it. The adaptive feature weighting module is used to recalibrate the initial feature representation based on a dual channel and spatial attention mechanism to enhance key features related to defects and suppress irrelevant background features, outputting a weighted feature vector. The key features include crack features, spalling features, or indentation features on the rail surface. The initial feature representation refers to the multi-channel feature map of the rail image extracted by the CNN backbone, used to characterize the original visual features related to defects in the rail image. Step 3: For each of the few sample recognition tasks, input the samples in the support set into the feature extraction network to obtain the weighted feature vector of each sample; based on the weighted feature vector of all samples of the same category, calculate the prototype feature vector of that category in the feature space; Step 4: Input the samples in the query set into the feature extraction network to obtain the weighted feature vector of each query sample; calculate the distance between the weighted feature vector of each query sample and the prototype feature vector of all categories obtained in Step 3. Step 5: Based on the distance, calculate the probability distribution of each query sample belonging to each category using the softmax function, and calculate the meta-training loss based on the difference between the probability distribution and the true label of the query sample; Step 6: Optimize the parameters of the feature extraction network using the gradient descent algorithm, minimize the meta-training loss, and iterate through multiple small sample recognition tasks to enable the feature extraction network to learn the ability to recognize new defect categories from a small number of samples, thus obtaining a trained feature extraction network. Step 7: In the meta-testing phase, the support set and query set containing the new defect categories are input into the trained feature extraction network to obtain the sample classification results in the query set containing the new defect categories.

[0027] Optionally, the adaptive feature weighting module is a channel-space dual attention module, which includes a channel attention submodule and a spatial attention submodule connected in sequence; The channel attention submodule is used to generate a weight vector based on the initial feature representation, and then multiply the weight vector with the initial feature representation channel by channel to generate a channel-weighted feature representation. The spatial attention submodule is used to generate a spatial dimension weight map based on the channel-weighted feature representation, and to multiply the channel-weighted feature representation position by position using the spatial dimension weight map to output the weighted feature vector.

[0028] Optionally, generating the weight vector includes: The initial feature representation is subjected to global average pooling and global max pooling respectively to obtain two channel descriptors; The two channel descriptors are input into a shared multilayer perceptron network, which includes a dimensionality reduction layer and a recovery layer. The two features output by the multilayer perceptron network are added element by element, and then passed through the sigmoid activation function to generate the weight vector.

[0029] Optionally, generating the weighted graph of the spatial dimension includes: The channel-weighted feature representations are subjected to average pooling and max pooling along the channel dimension to obtain two spatial feature maps. The two spatial feature maps are concatenated along the channel dimension; The concatenated feature maps are then subjected to convolution operations and then activated by the Sigmoid function to generate a weight map with spatial dimensions.

[0030] Optionally, in step 3, the formula for calculating the prototype feature vector of the nth type of defect in the feature space is:

[0031] in, Let n be the prototype feature vector of the nth type of defect in the feature space. Indicates the class n in the support set. One labeled sample, express The weighted feature vector obtained after passing through the feature extraction network has a value of n ranging from 1 to N, and K is the total number of samples of the nth type of defect.

[0032] Optionally, in step 5, the calculated distance is either Euclidean distance or the inverse of cosine similarity; when using Euclidean distance, the probability that the query sample x belongs to the nth type of defect is... The calculation formula, expressed using the softmax function, is as follows:

[0033] Where y is the class label variable of the sample; Given model parameters θ, the probability that the category label y of the query sample x is the nth type of defect; The weighted feature vector obtained after the query sample x has passed through the feature extraction network is... Let n be the prototype feature vector of the nth type of defect in the feature space. express and The Euclidean distance between them; The summation index variable has a range of 1 to N, and exp is an exponential function.

[0034] Optionally, the meta-training loss function in step 5 is the negative log-likelihood loss, and its calculation formula is as follows:

[0035] Where T is the total number of samples in the query set. This represents the t-th sample in the query set; for The true label, For query samples Belongs to real labels The probability of the category, This is the training loss.

[0036] This application also provides an embodiment of a rail small-sample defect identification system for implementing the method described in the above embodiment, comprising: The task construction module is used to construct a meta-training task set, which consists of multiple few-sample recognition tasks. These tasks are generated iteratively from a meta-training set containing images of multiple types of rail defects. Each few-sample recognition task includes a support set and a query set. The support set contains N types of defect samples, each type of defect sample having K labeled samples. The query set contains labeled samples that are not repeated from the labeled samples in the support set, and each sample in the query set belongs to one of the N types of defects in the support set. A network construction module is used to construct a feature extraction network, which includes a convolutional neural network (CNN) backbone and an adaptive feature weighting module connected to the CNN backbone. The CNN backbone is used to extract an initial feature representation of the rail image input to it. The adaptive feature weighting module is used to recalibrate the initial feature representation based on a dual channel and spatial attention mechanism to enhance key features related to defects and suppress irrelevant background features, outputting a weighted feature vector. The key features include crack features, spalling features, or indentation features on the rail surface. The initial feature representation refers to the multi-channel feature map of the rail image extracted by the CNN backbone, used to characterize the original visual features related to defects in the rail image. The feature vector calculation module is used to input the samples in the support set into the feature extraction network for each of the few sample recognition tasks to obtain the weighted feature vector of each sample; and to calculate the prototype feature vector of the category in the feature space based on the weighted feature vector of all samples of the same category. The distance calculation module is used to input the samples in the query set into the feature extraction network to obtain the weighted feature vector of each query sample; and to calculate the distance between the weighted feature vector of each query sample and the prototype feature vector of all categories. The meta-training loss calculation module is used to calculate the probability distribution of each query sample belonging to each category based on the distance using the softmax function, and to calculate the meta-training loss based on the difference between the probability distribution and the true label of the query sample. The model training module is used to optimize the parameters of the feature extraction network through the gradient descent algorithm, minimize the meta-training loss, and through multiple iterations of the few sample recognition tasks, enable the feature extraction network to learn the ability to recognize new defect categories from a small number of samples, thus obtaining a trained feature extraction network. The classification and recognition module is used in the meta-testing stage to input the support set and query set containing the new defect category into the trained feature extraction network to obtain the sample classification results in the query set containing the new defect category.

[0037] This application also provides an embodiment of a computer-readable storage medium having a computer program stored thereon, which, when executed by a processor, implements the methods described in the above embodiments.

[0038] This application also provides an embodiment of an electronic device, including a memory, a processor, and a computer program stored in the memory and executable on the processor, wherein the processor executes the computer program to implement the methods described in the above embodiments.

[0039] The core of this application's embodiments lies in constructing a meta-learning training framework. This framework is trained by simulating real-world few-shot recognition scenarios: each training iteration (or episode) randomly samples from the dataset to construct a few-shot recognition task containing a support set and a query set. The model (feature extraction network) gains the ability to generalize core features of new categories from a small number of samples through repeated learning on such tasks.

[0040] To overcome the challenges posed by the complex background and subtle defect areas in rail images, this application embeds an adaptive feature weighting module into the feature extraction network. This module, through a collaborative channel attention and spatial attention mechanism, dynamically evaluates and enhances key feature channels and spatial regions related to defects, while suppressing interference from irrelevant background information such as ballast, fasteners, and uneven lighting, thereby extracting more discriminative feature representations from the source.

[0041] In the meta-learning process, for each task, features are first extracted using a small number of samples from the support set through the aforementioned network, and the average feature vector of each category, i.e., the "prototype," is calculated. Subsequently, by measuring the distance between the query set sample features and the prototypes of each category in the feature space (e.g., Euclidean distance), and using the Softmax function to transform them into a probability distribution, small-sample classification of the query samples can be achieved. This approach addresses the fundamental problem of "how to extract pure, highly discriminative feature representations from finite samples," effectively improving the performance of standard prototype networks in complex industrial scenarios.

[0042] The solution provided in this application not only includes the aforementioned algorithm model, but also covers the corresponding system architecture, and can be integrated into various hardware carriers such as inspection robots, vehicle-mounted intelligent devices, or cloud analysis platforms, providing a feasible technical path for the engineering deployment, industrial application, and network-level collaborative operation and maintenance of intelligent rail defect detection technology.

[0043] The above embodiments are merely examples. Based on the technical solutions of this application, those skilled in the art can make various equivalent modifications or substitutions without creative effort, according to the disclosed technical content. For example, the convolutional neural network backbone can be replaced with other mainstream network structures such as ResNet-34 and DenseNet; the attention enhancement mechanism can also be combined with other meta-learning frameworks (such as MAML); furthermore, the method provided in this application is also applicable to small-sample defect identification scenarios for other industrial components such as bearings and blades. These variations or applications should all be included within the protection scope of this application.

[0044] Example 1: A rail defect identification method based on prototype network and dual attention mechanism: This embodiment illustrates in detail the application of the method of this application in a specific experimental scenario, aiming to demonstrate its feasibility and superiority.

[0045] (1) Data preparation and preprocessing: To construct the meta-training task set, this embodiment prepares the following basic data: Basic dataset: The experiment uses a self-built image database of rail surface defects. This database contains images of a total of P=9 common defect categories, including cracks, spalling, shelling, craze, and corrosion. Each defect category contains approximately M=100 high-quality sample images annotated by experts.

[0046] Dataset partitioning: To adhere to the meta-learning evaluation paradigm, the entire dataset is randomly divided into a meta-training set, a meta-validation set, and a meta-test set in a 7:2:1 ratio. The defect categories within these three subsets are mutually exclusive; that is, defect categories in the meta-test set never appear during the meta-training and meta-validation phases. This rigorously evaluates the feature extraction network's ability to generalize to new defect categories.

[0047] Image preprocessing: To adapt to the input requirements of the convolutional neural network backbone in the feature extraction network, all sample images were uniformly scaled to 224×224 pixels. Then, standardization was performed: for each pixel's RGB channels, the mean of the ImageNet dataset [0.485, 0.456, 0.409] was subtracted, and then divided by the standard deviation [0.229, 0.224, 0.225]. During the meta-training phase, to improve model robustness, random horizontal flipping was applied to the input images as a basic data augmentation technique; during the meta-testing phase, no data augmentation was used.

[0048] (2) Model building: The specific structure of the feature extraction network constructed in this embodiment is as follows: Figure 2 As shown, it includes a convolutional neural network backbone and an adaptive feature weighting module connected to it, and its specific structure is as follows: ① Backbone Network: ResNet-18, pre-trained on the ImageNet dataset, was selected as the backbone of the convolutional neural network. The global average pooling layer and fully connected classification layer at the end of its original architecture were removed, and its convolutional parts were retained. The output of the last convolutional group (i.e., layer 4) was taken as the initial feature representation F, with dimensions [512, 7, 7] (number of channels C × height H × width W).

[0049] ② Adaptive Feature Weighting Module: This module is a channel-spatial dual attention module, cascaded after the backbone network, used to recalibrate the initial feature representation based on the channel and spatial dual attention mechanism. Its detailed structure is as follows... Figure 2 As shown, it includes a channel attention submodule and a spatial attention submodule connected in sequence: Channel attention submodule: Input: Initial feature representation of the backbone network output .

[0050] Operation: Global average pooling (GAP) and global max pooling (GMP) are performed on F to obtain two 512-dimensional channel-level feature descriptors. These two descriptors are then fed into a shared two-layer fully connected network (MLP). The first layer of this MLP reduces the dimension from 512 to 512 / 16 = 32 (using the ReLU activation function), and the second layer restores the dimension from 32 to 512.

[0051] Weight generation: The two feature vectors processed by the MLP are added element-wise, and then the final channel weight vector is generated by passing the sigmoid activation function. .

[0052] Output: The generated channel weight vector The channel-weighted feature representation is obtained by multiplying the input feature F channel by channel. .

[0053] Spatial Attention Submodule: Input: Channel-weighted feature map .

[0054] Operation: To Average pooling and max pooling are performed along the channel dimension to obtain two feature maps with a spatial size of 7×7. These two feature maps are then concatenated along the channel dimension to form a 2-channel [2,7,7] feature map.

[0055] Weight generation: A 7×7 convolutional layer is used to process the concatenated feature map, compressing it into a single-channel [1,7,7] feature map. Then, a sigmoid activation function is applied to generate a spatial weight map. .

[0056] Output: Input features Element-wise multiplication yields the final weighted feature vector enhanced by dual attention. .

[0057] (3) Meta-training process: The goal of meta-training is to enable the feature extraction network to learn the "meta-capabilities" of new categories from a small number of samples. The training process is conducted in "episodes" (task batches), with each iteration corresponding to a complete small sample task. The process is as follows: Figure 1 As shown in the "meta-training phase".

[0058] Task Sampling: In each iteration, N=5 different defect categories are randomly selected from the meta-training set, and K=5 labeled samples are randomly selected from each category to form the support set. Simultaneously, Q=15 different labeled samples are randomly selected from each of these 5 categories to form the query set, ensuring that the selected samples are distinct from those in the support set. This constructs a "5-way 5-shot" small sample recognition task, constituting one task within the meta-training task set.

[0059] like Figure 3 As shown in the figure, taking N=3, K=2, Q=3 as an example, the construction process of a small sample identification task is illustrated: First, starting from a sample containing P defect categories ( Figure 3 In the example P, there are multiple classes, such as class A, class B, class C, class D... class N. N different defect categories are randomly selected from the rail defect image database (i.e., the meta-training set). Figure 3 The example extracts classes A, C, and E; then, from these N classes, K labeled samples are randomly selected from each class to form a support set. Figure 3 The example extracts from class A and Extract from class C and Extract from class E and Finally, from the same N categories, Q labeled samples that are distinct from the support set samples are randomly selected from each category to form the query set. Figure 3 The example extracts from class A , , Extract from class C , , Extract from class E , , Where P represents the total number of defect categories in the rail defect image database, and M represents the total number of samples in each category (). Figure 3 (Not explicitly stated, but implied to be M samples per class), N represents the number of defect categories extracted in each small sample identification task, K represents the number of samples for each category in the support set, and Q represents the number of samples for each category in the query set.

[0060] Prototype computation: All N×K=25 images from the support set are input into the feature extraction network to obtain a weighted feature vector for each image, with a dimension of 512. For each extracted class n (n=1 to 5), the arithmetic mean of the weighted feature vectors of its corresponding K=5 support samples is taken to obtain the prototype feature vector for that class. .

[0061] Distance metric and classification: The N×Q=75 images of the query set are input into the same network to obtain their feature vectors. For each query sample, a weighted feature vector is generated. For each query feature vector, calculate its Euclidean distance to each of the N prototype feature vectors. Using this distance as a logistic value (logits), the Softmax function is used to calculate the probability distribution of each sample in the query set belonging to each category. .

[0062] Loss Calculation and Optimization: The meta-training loss is used as the objective function, and negative log-likelihood loss is employed for calculation. The Adam optimizer is used for parameter optimization to minimize the meta-training loss; the initial learning rate is set to 1e-3, and a step-decay strategy is adopted (e.g., the learning rate is halved every 500 episodes). A total of 8000 episodes are trained, and performance is monitored on the meta-validation set to select the best model. After iterations on multiple few-sample recognition tasks, the feature extraction network acquires the ability to learn new defect categories from a small number of samples, resulting in a well-trained feature extraction network.

[0063] (4) Meta-test reasoning process: After training is complete, all parameters of the feature extraction network are fixed, and the model enters the meta-testing stage to evaluate its ability to generalize to new defect categories.

[0064] From the meta-test set (which includes defect categories that did not appear during the training phase), repeatedly random samples are used to construct 500 “5-way K-shot” test tasks (K can be 1 or 5).

[0065] For each test task, no parameter updates are performed. The process is the same as during training: the prototype feature vector of the new category is calculated using a small number of samples from the support set, and then the query set samples are classified to obtain the classification results of the samples in the query set that contain the new defect category.

[0066] The average classification accuracy and its 95% confidence interval for all query samples in these 500 test tasks were statistically analyzed and used as the final performance metric of the model. This accuracy was derived based on a probability calculation formula.

[0067] (5) Effect verification and comparison experiment: To quantitatively verify the effectiveness of the method in this application, the following comparative experiment was designed.

[0068] ① Experimental setup: Under the same meta-training set, meta-validation set, and meta-test set partitioning and evaluation protocol, the following three methods were compared: Baseline 1 (Finetune): Uses a convolutional neural network backbone (ResNet-18), but removes the adaptive feature weighting module. On the support set of the test task, only the last fully connected layer of the network is fine-tuned, and then tested on the query set.

[0069] Baseline 2 (ProtoNet): Standard prototype network. It uses the same convolutional neural network backbone (ResNet-18) as in this application, but without adding any adaptive feature weighting modules. It directly uses the initial feature representations extracted by the backbone network to compute the prototype and classification.

[0070] The method of this application adopts the complete feature extraction network structure described in this embodiment, namely the combination of a convolutional neural network backbone (ResNet-18) and a channel-space dual attention module.

[0071] ② Evaluation metrics and results: Evaluation was conducted under two more challenging settings: 5-way 1-shot (only 1 supporting sample per class) and 5-way 5-shot (5 supporting samples per class). Statistical results after running 500 random tasks on the independent meta-test set are shown in Table 1.

[0072] Table 1 Statistical Results

[0073] ③ Results Analysis: The proposed method significantly outperforms the two baseline methods in both small-sample settings. In particular, in the extremely scarce 1-shot scenario, it achieves a significant improvement of more than 8 percentage points compared to the standard ProtoNet, effectively verifying the crucial role of the adaptive feature weighting module in focusing on key defect features, suppressing background interference, and thus improving the discriminative power of the prototype feature vector.

[0074] Ablation experiments show that the accuracy of 5-way 5-shot is 82.1% when using only the channel attention submodule and 80.9% when using only the spatial attention submodule. Both are lower than the combined effect (86.8%), demonstrating that the channel and spatial attention mechanisms are complementary in rail defect identification tasks and achieve the best results when working together.

[0075] The 86.8% 5-way 5-shot accuracy rate has significant practical value. This means that during on-site inspections, only five labeled samples of a certain novel defect are needed to form a support set, and the system's accuracy in identifying this type of defect can reach a practical level. This greatly lowers the threshold for data collection and labeling, making rapid response to unknown defects possible.

[0076] Example 2: System integrated into a mobile rail inspection robot: This embodiment describes the deployment of the feature extraction network trained in this application into an actual hardware system: an autonomous mobile rail inspection robot, forming an end-to-end solution. This system itself constitutes a neural network classifier recognition system.

[0077] (1) System composition: Image acquisition unit: A vision acquisition system is mounted on the front of the robot body, including a high-resolution industrial area array camera and a set of strip LED supplementary lights. As the robot moves autonomously along the track, the system continuously and stably acquires images of the rail surface at a fixed frame rate (e.g., 10fps).

[0078] Computational processing unit: A high-performance embedded AI computing device (such as NVIDIA Jetson AGX Xavier) is integrated inside the robot. This device is pre-loaded with a feature extraction network and a prototype feature vector library corresponding to known defect categories, forming the "brain" of the system for executing the methods of this application.

[0079] Control and Communication Unit: Responsible for the robot's motion control and task scheduling, and transmitting detection results and alarm information back to the back-end operation and maintenance management center in real time via 4G / 5G or Wi-Fi.

[0080] (2) Workflow ("Detection and Learning" mode) Routine Inspection: The robot inspects along a predetermined route, and the system analyzes each frame of the captured image in real time. For each frame, a weighted feature vector is first extracted using a feature extraction network. Then, the Euclidean distance between this vector and the prototype feature vector of the known defect category is calculated, and the classification probability is obtained through a softmax function, thereby achieving rapid identification of known defects. Upon discovery of a defect, it is immediately marked, located, and an alarm is triggered.

[0081] Learning new defects from small samples: When inspectors discover a suspected defect that has not been entered into the system through the background video stream, they can remotely control the robot to pause and clearly extract K typical images (e.g., K=5) of the defect from the video stream. These images are then assigned a temporary label (e.g., "Unknown crack - section X") through the human-computer interaction interface. These K images constitute a support set for a new category, and all of these images are already labeled samples.

[0082] Online inference: The computational processing unit then initiates a small-sample inference process, which inputs K images from the support set into the feature extraction network to obtain a weighted feature vector for each image; calculates the arithmetic mean of these K weighted feature vectors as the prototype feature vector for the new defect category; and adds the new prototype to the runtime prototype feature vector library.

[0083] Subsequently, the system continues to inspect and extracts the weighted feature vectors of the subsequently acquired images in real time, and performs distance measurement and classification with the updated prototype library (which includes all categories of prototypes, both old and new).

[0084] Results feedback: When the system identifies the new type of defect again with high confidence, it will mark it on the video stream in real time and record the precise location coordinates of the defect (obtained by fusion positioning of robot odometry and GPS) to form a complete defect report.

[0085] This system achieves an online small-sample learning closed loop of "on-site discovery, on-site learning, and on-site recognition." Inspection personnel no longer need to send data back to the center for days or weeks of retraining; the robot can be taught to recognize a new defect within minutes on-site, greatly improving the flexibility, timeliness, and intelligence of inspections. It reduces absolute dependence on central servers and stable networks, making it suitable for deployment in remote areas with poor network conditions, and enhancing the system's autonomy and adaptability.

[0086] Example 3: A cloud-based collaborative rail defect diagnosis platform: This embodiment describes a more generalized, system-level application architecture, demonstrating the collaborative value of the method in road network-level operation and maintenance. The platform itself can be considered an extended implementation of the system described in this application.

[0087] (1) System architecture: Edge network (each railway bureau / engineering section): Deploy a lightweight, high real-time feature extraction network (a lightweight version of the feature extraction network in this application, or a simplified version containing only its convolutional neural network backbone). It is responsible for daily high-frequency inspections and preliminary screenings within its jurisdiction, identifying most known defects. When a suspected new defect with low confidence or not found in the prototype feature vector library is discovered, it automatically encrypts a small number of clear samples (i.e., labeled samples constituting the support set), environmental context, and location information, and uploads them to the cloud center.

[0088] The cloud center runs a continuously evolving, powerful feature extraction network (i.e., the core model trained according to the method of this application). It aggregates newly labeled defect samples uploaded from all edge terminals of the entire road network, automatically constructs incremental few-sample recognition tasks, and continuously performs meta-training and fine-tuning on the feature extraction network, so that its prototype feature vector library is constantly enriched and optimized, enabling it to recognize an increasing number of new defect types.

[0089] Model distribution mechanism: The cloud center regularly (e.g., weekly) or according to the urgency of the task, securely distributes the optimized and updated feature extraction network parameters, or the prototype feature vectors of newly added defect categories, to each edge device.

[0090] (2) Workflow example: During a routine inspection, a new type of "stress fatigue cracking" defect was discovered for the first time, which the on-site equipment could not identify. Staff used a terminal app to label five typical images as "new type - cracking" and uploaded them to the cloud. These five images constitute the support set for this new category, and all are labeled samples.

[0091] Upon receiving the images, the cloud center utilizes these five images and powerful computing resources to complete a targeted small-sample recognition task within minutes: inputting the support set images into the cloud feature extraction network to obtain the weighted feature vector of each image, and calculating its arithmetic mean as the prototype feature vector of the new defect category, thereby quickly integrating the recognition capability of the new defect into the cloud main model.

[0092] After the update command is issued, all detection equipment (including inspection robots, vehicle-mounted equipment, and fixed monitoring points) deployed with the system on the entire road network will immediately obtain the updated prototype feature vector library during the next synchronization, thus enabling them to identify this type of "stress fatigue cracking".

[0093] This embodiment establishes a collaborative evolution ecosystem for rail defect knowledge based on "distributed discovery, centralized learning, and networked sharing." It enables the real-time sharing and synchronization of operation and maintenance experience and knowledge across the entire railway network; a single discovery benefits the entire network. This significantly amplifies the value of the application's rapid adaptive capability with small samples, achieving continuous and collaborative evolution of the overall intelligence level of the detection system.

[0094] See Figure 1 It clearly demonstrates the core steps, data flow, and inter-module interactions of the method described in this application during the meta-training and meta-testing phases. It fully embodies the complete technical chain of "task sampling → feature extraction and adaptive enhancement → prototype computation → distance measurement and classification → meta-optimization / inference" and serves as the overall framework for understanding the methodology of this application.

[0095] See Figure 2 This paper details how the adaptive feature weighting module (i.e., the channel-space dual attention module) described in this application is composed of a channel attention submodule and a spatial attention submodule connected in series. It includes the key operations of each submodule (such as global pooling, shared multilayer perceptron, and convolutional layer) and the transformation process of feature data. This is the key to understanding how this application achieves the core function of "enhancing key features and suppressing background interference".

[0096] See Figure 3 This paper illustrates in an intuitive way how a small sample recognition task instance for meta-training or meta-testing can be constructed by randomly sampling N categories from a meta-training set containing multiple types of rail defects, and then sampling K labeled samples from each category to form a support set and Q different labeled samples to form a query set.

[0097] Table 2 shows a performance comparison of the method in this application with existing traditional methods on small-sample rail defect identification tasks. It quantitatively compares the average classification accuracy and confidence interval of different methods (such as support set fine-tuning, standard prototype network, and the method described in this application) under the same 5-way 1-shot and 5-way 5-shot test settings, and uses objective data to intuitively demonstrate the effectiveness and superiority of the method in this application.

[0098] Table 2 Performance Comparison

[0099] Finally, it should be noted that the above embodiments are only used to illustrate the technical solutions of this application and not to limit them. Although this application has been described in detail with reference to the above embodiments, those skilled in the art should understand that modifications or equivalent substitutions can still be made to the specific implementation of this application. Any modifications or equivalent substitutions that do not depart from the spirit and scope of this application should be covered within the protection scope of the claims of this application.

Claims

1. A method for identifying defects in small samples of steel rails, characterized in that, include: Step 1: Construct a meta-training task set, which consists of multiple few-shot recognition tasks. Each few-shot recognition task is constructed as follows: From the meta-training set containing images of multiple types of rail defects, randomly select N different defect categories; for each of the N selected defect categories, randomly select K labeled samples from each defect category to form a support set; simultaneously, from the same N defect categories, randomly select Q labeled samples for each category, and the selected samples are not repeated with the samples in the support set to form a query set. Step 2: Construct a feature extraction network, which includes a convolutional neural network (CNN) backbone and an adaptive feature weighting module connected to the CNN backbone. The CNN backbone is used to extract the initial feature representation of the rail image input to it. The adaptive feature weighting module is used to recalibrate the initial feature representation based on a dual channel and spatial attention mechanism to enhance key features related to defects and suppress irrelevant background features, outputting a weighted feature vector. The key features include crack features, spalling features, or indentation features on the rail surface. The initial feature representation refers to the multi-channel feature map of the rail image extracted by the CNN backbone, used to characterize the original visual features related to defects in the rail image. Step 3: For each of the few sample recognition tasks, input the samples in the support set into the feature extraction network to obtain the weighted feature vector of each sample; based on the weighted feature vector of all samples of the same category, calculate the prototype feature vector of that category in the feature space; Step 4: Input the samples in the query set into the feature extraction network to obtain the weighted feature vector of each query sample; calculate the distance between the weighted feature vector of each query sample and the prototype feature vector of all categories obtained in Step 3. Step 5: Based on the distance, calculate the probability distribution of each query sample belonging to each category using the softmax function, and calculate the meta-training loss based on the difference between the probability distribution and the true label of the query sample; Step 6: Optimize the parameters of the feature extraction network using the gradient descent algorithm, minimize the meta-training loss, and iterate through multiple small sample recognition tasks to enable the feature extraction network to learn the ability to recognize new defect categories from a small number of samples, thus obtaining a trained feature extraction network. Step 7: In the meta-testing phase, the support set and query set containing the new defect categories are input into the trained feature extraction network to obtain the sample classification results in the query set containing the new defect categories.

2. The method according to claim 1, characterized in that, The adaptive feature weighting module is a channel-space dual attention module, which includes a channel attention submodule and a spatial attention submodule connected in sequence; The channel attention submodule is used to generate a weight vector based on the initial feature representation, and then multiply the weight vector with the initial feature representation channel by channel to generate a channel-weighted feature representation. The spatial attention submodule is used to generate a spatial dimension weight map based on the channel-weighted feature representation, and to multiply the channel-weighted feature representation position by position using the spatial dimension weight map to output a weighted feature vector.

3. The method according to claim 2, characterized in that, The generation of the weight vector includes: The initial feature representation is subjected to global average pooling and global max pooling respectively to obtain two channel descriptors; The two channel descriptors are input into a shared multilayer perceptron network, which includes a dimensionality reduction layer and a recovery layer. The two features output by the multilayer perceptron network are added element by element, and then passed through the sigmoid activation function to generate the weight vector.

4. The method according to claim 2, characterized in that, The generation of the spatial dimension weighted graph includes: The channel-weighted feature representations are subjected to average pooling and max pooling along the channel dimension to obtain two spatial feature maps. The two spatial feature maps are concatenated along the channel dimension; The concatenated feature maps are then subjected to convolution operations and then activated by the Sigmoid function to generate a weight map with spatial dimensions.

5. The method according to claim 1, characterized in that, In step 3, the formula for calculating the prototype feature vector of the nth type of defect in the feature space is: in, Let n be the prototype feature vector of the nth type of defect in the feature space. Indicates the class n in the support set. One labeled sample, express The weighted feature vector obtained after passing through the feature extraction network has a value of n ranging from 1 to N, and K is the total number of samples of the nth type of defect.

6. The method according to claim 5, characterized in that, In step 5, the calculated distance is either the Euclidean distance or the inverse of the cosine similarity; when using the Euclidean distance, the probability that the query sample x belongs to the nth type of defect is calculated. The calculation formula, expressed using the softmax function, is as follows: Where y is the class label variable of the sample; Given model parameters θ, the probability that the category label y of the query sample x is the nth type of defect; The weighted feature vector obtained after the query sample x has passed through the feature extraction network is... Let n be the prototype feature vector of the nth type of defect in the feature space. express and The Euclidean distance between them; The summation index variable has a value range of 1 to N, and exp is an exponential function.

7. The method according to claim 6, characterized in that, The meta-training loss function in step 5 is the negative log-likelihood loss, and its calculation formula is as follows: Where T is the total number of samples in the query set. This represents the t-th sample in the query set; for The true label, For query samples Belongs to real labels The probability of the category, This is the training loss.

8. A rail small-sample defect identification system, used to implement the method as described in any one of claims 1-7, characterized in that, include: The task construction module is used to construct a meta-training task set, which consists of multiple few-sample recognition tasks. These tasks are generated iteratively from a meta-training set containing images of multiple types of rail defects. Each few-sample recognition task includes a support set and a query set. The support set contains N types of defect samples, each type of defect sample having K labeled samples. The query set contains labeled samples that are not repeated from the labeled samples in the support set, and each sample in the query set belongs to one of the N types of defects in the support set. A network construction module is used to construct a feature extraction network, which includes a convolutional neural network (CNN) backbone and an adaptive feature weighting module connected to the CNN backbone. The CNN backbone is used to extract an initial feature representation of the rail image input to it. The adaptive feature weighting module is used to recalibrate the initial feature representation based on a dual channel and spatial attention mechanism to enhance key features related to defects and suppress irrelevant background features, outputting a weighted feature vector. The key features include crack features, spalling features, or indentation features on the rail surface. The initial feature representation refers to the multi-channel feature map of the rail image extracted by the CNN backbone, used to characterize the original visual features related to defects in the rail image. The feature vector calculation module is used to input the samples in the support set into the feature extraction network for each of the few sample recognition tasks to obtain the weighted feature vector of each sample; and to calculate the prototype feature vector of the category in the feature space based on the weighted feature vector of all samples of the same category. The distance calculation module is used to input the samples in the query set into the feature extraction network to obtain the weighted feature vector of each query sample; and to calculate the distance between the weighted feature vector of each query sample and the prototype feature vector of all categories. The meta-training loss calculation module is used to calculate the probability distribution of each query sample belonging to each category based on the distance using the softmax function, and to calculate the meta-training loss based on the difference between the probability distribution and the true label of the query sample. The model training module is used to optimize the parameters of the feature extraction network through the gradient descent algorithm, minimize the meta-training loss, and through multiple iterations of the few sample recognition tasks, enable the feature extraction network to learn the ability to recognize new defect categories from a small number of samples, thus obtaining a trained feature extraction network. The classification and recognition module is used in the meta-testing stage to input the support set and query set containing the new defect category into the trained feature extraction network to obtain the sample classification results in the query set containing the new defect category.

9. A computer-readable storage medium having a computer program stored thereon, characterized in that, When the computer program is executed by a processor, it implements the method as described in any one of claims 1-7.

10. An electronic device comprising a memory, a processor, and a computer program stored in the memory and executable on the processor, characterized in that, When the processor executes the computer program, it implements the method as described in any one of claims 1-7.

Images