Rail surface defect segmentation method and device based on ordered cross-scale feature interaction

Abstract

The present invention relates to the technical field of rail surface defect segmentation, and provides a rail surface defect segmentation method and device based on ordered cross-scale feature interaction. The rail surface defect segmentation model sequentially realizes global dependent feature extraction, cross-scale feature channel-level interaction, feature space-level interaction, and feature pixel-level interaction, which can powerfully express the defects on the rail surface, fully enhance the features, highlight the defect areas in the features, and suppress irrelevant background areas. It is more suitable for defect features with complex backgrounds, low foreground-background contrast, and blurred target edges. The obtained features are more conducive to the subsequent defect segmentation process. Furthermore, by using the target features, a more accurate and detailed surface defect segmentation map of the rail surface image can be obtained, and accurate pixel-level positioning of the rail surface defects can be achieved.

Description

Technical Field

[0001] The present invention relates to the technical field of rail surface defect segmentation, and in particular to a rail surface defect segmentation method and device based on ordered cross-scale feature interaction. Background Art

[0002] The safety of railway transportation is of vital importance. For railway transportation, the inspection of rails on the track is very important for railway safety. Similarly, the manufacturing process and factory quality inspection of off-rail rails are also very important. The segmentation of rail surface defects can not only quickly locate the defects of on-rail tracks and avoid potential accidents, but also improve the yield rate of the manufacturing process of off-rail rails and ensure the quality of off-rail rails leaving the production workshop. Rail surface defect images are weak semantic images. Compared with natural images, they have the characteristics of complex background, low foreground and background contrast, and blurred target edges. This also leads to the development of the industry. Traditional rail surface defect segmentation solutions, including manual segmentation solutions, automated segmentation solutions, and traditional machine learning solutions, have been difficult to meet the needs of inspection and factory production. Therefore, how to achieve high-precision rail surface defect segmentation in complex environments is still one of the industry's pain points.

[0003] Existing rail surface defect segmentation methods can be divided into two categories, namely, methods based on manually designed features and methods based on deep learning. Methods based on manually designed features are often based on theories such as information theory, control theory, sparse representation, and frequency domain analysis. The extracted features are not universal and cannot accurately express defects, which limits the defect segmentation performance. Methods based on deep learning learn from a large amount of data and can learn features that are conducive to defect segmentation, significantly improving the performance of defect segmentation. However, methods based on deep learning usually use convolutional neural networks as the backbone to extract features, which cannot fully model the relationship between complex defect areas and backgrounds, and cannot extract global dependent features. In addition, methods based on deep learning often use relatively simple methods to enhance the extracted features, which cannot effectively highlight the defect areas in the features. At the same time, rail surface defect images are very rare, and data collection is very time-consuming. Summary of the invention

[0004] The present invention provides a rail surface defect segmentation method and device based on ordered cross-scale feature interaction, which are used to solve the defects existing in the prior art.

[0005] The present invention provides a rail surface defect segmentation method based on ordered cross-scale feature interaction, comprising:

[0006] Acquire rail surface images;

[0007] Based on the rail surface defect segmentation model, globally dependent features are extracted from the rail surface image to obtain multiple global dependent features of different scales, and cross-scale feature channel-level interaction, cross-scale feature space-level interaction, and cross-scale feature pixel-level interaction are sequentially performed on the multiple global dependent features to obtain target features, and a surface defect segmentation map of the rail surface image is determined based on the target features;

[0008] The cross-scale feature channel-level interaction is to sequentially perform channel splitting, channel interaction fusion, and channel splicing on two adjacent global dependency features to obtain two first enhancement results;

[0009] The cross-scale feature space-level interaction is to normalize the size of the two first enhancement results, and then perform length splitting, spatial interaction fusion and length splicing in sequence to obtain two second enhancement results;

[0010] The cross-scale feature pixel-level interaction is to normalize the size of the two second enhancement results, and then perform feature interaction fusion based on the channel-level self-attention mechanism to obtain two third enhancement results, and all third enhancement results constitute the target feature.

[0011] According to a rail surface defect segmentation method based on ordered cross-scale feature interaction provided by the present invention, the cross-scale feature channel-level interaction specifically includes:

[0012] Perform channel splitting on the two global dependency features respectively to obtain two first sub-feature sets;

[0013] The first sub-features in the two first sub-feature sets are grouped accordingly to obtain multiple first sub-feature groups, and the two first sub-features in each first sub-feature group are subjected to spatial global average pooling and spatial global maximum pooling, and the global average pooling results and the global maximum pooling results of the two first sub-features are subjected to channel splicing, channel shuffling, two full connection operations, pixel-by-pixel summation, activation, and channel splitting in sequence to obtain first split features and second split features, and the first split features and the second split features are respectively multiplied pixel by pixel with the corresponding first sub-features to obtain two first enhanced features;

[0014] The two first enhancement features corresponding to each first sub-feature group are channel-joined in sequence to obtain the two first enhancement results.

[0015] According to a rail surface defect segmentation method based on ordered cross-scale feature interaction provided by the present invention, the cross-scale feature space-level interaction specifically includes:

[0016] After normalizing the sizes of the two first enhancement results respectively, splitting them in the length dimension to obtain two second sub-feature sets;

[0017] The second sub-features in the two second sub-feature sets are grouped accordingly to obtain multiple second sub-feature groups, and the two second sub-features in each second sub-feature group are spatially spliced ​​and then channel-level global average pooling and channel-level global maximum pooling are performed, and the global average pooling results and the global maximum pooling results of the two second sub-features are sequentially subjected to channel splicing, convolution operation, activation, and spatial splitting to obtain a third split feature and a fourth split feature, and the third split feature and the fourth split feature are respectively pixel-by-pixel multiplied with the corresponding second sub-features to obtain two second enhanced features;

[0018] The two second enhancement features corresponding to each second sub-feature group are respectively spliced ​​in length in sequence to obtain the two second enhancement results.

[0019] According to a rail surface defect segmentation method based on ordered cross-scale feature interaction provided by the present invention, the cross-scale feature pixel-level interaction specifically includes:

[0020] Normalizing the sizes of the two second enhancement results respectively to obtain a third feature and a fourth feature;

[0021] After downsampling the third feature, two convolution operations are respectively performed to obtain a first value matrix and a query matrix, and two convolution operations are respectively performed to obtain a key matrix and a second value matrix on the fourth feature;

[0022] Transposing the query matrix to obtain a transposed matrix, and calculating an attention score based on the transposed matrix and the key matrix;

[0023] Transpose the attention score and perform matrix multiplication with the first value matrix after passing a normalized exponential function to obtain a fifth feature, and perform matrix multiplication with the second value matrix after passing a normalized exponential function to obtain a sixth feature;

[0024] Upsampling the fifth feature to obtain an upsampling result, and summing the upsampling result with the third feature to obtain a first summation result, and summing the sixth feature with the fourth feature to obtain a second summation result;

[0025] The first summation result and the second summation result are respectively subjected to convolution operations to obtain the two third enhancement results.

[0026] According to a rail surface defect segmentation method based on ordered cross-scale feature interaction provided by the present invention, the surface defect segmentation map of the rail surface image is determined based on the target feature, including:

[0027] Based on the surface defect reasoning network in the rail surface defect segmentation model, the target feature is processed to obtain the surface defect segmentation map;

[0028] The surface defect reasoning network includes two convolutional layers, a random inactivation layer and a deconvolution layer connected in sequence.

[0029] According to a rail surface defect segmentation method based on ordered cross-scale feature interaction provided by the present invention, the rail surface defect segmentation model is based on which global dependent features are extracted from the rail surface image to obtain multiple global dependent features of different scales, including:

[0030] Based on the PVT network in the rail surface defect segmentation model, globally dependent features are extracted from the rail surface image to obtain four initial features of different scales;

[0031] The number of channels of the four initial features are unified to obtain four global dependency features.

[0032] According to a rail surface defect segmentation method based on ordered cross-scale feature interaction provided by the present invention, the rail surface defect segmentation model is trained based on the following steps:

[0033] Determine a training sample set, wherein the training sample set includes rail surface defect image samples and their corresponding segmentation truth images;

[0034] Inputting the rail surface defect image sample into an initial segmentation model to obtain a segmentation result output by the initial segmentation model;

[0035] Based on the segmentation result and the segmentation truth map, a binary cross entropy loss and an intersection-over-union loss are calculated, and based on the binary cross entropy loss and the intersection-over-union loss, a total loss is calculated;

[0036] Based on the total loss, the structural parameters of the initial segmentation model are iteratively updated using a gradient descent method to obtain the rail surface defect segmentation model.

[0037] The present invention also provides a rail surface defect segmentation device based on ordered cross-scale feature interaction, comprising:

[0038] An image acquisition module, used for acquiring rail surface images;

[0039] A defect segmentation module is used to extract global dependent features of the rail surface image based on the rail surface defect segmentation model to obtain multiple global dependent features of different scales, and sequentially perform cross-scale feature channel-level interaction, cross-scale feature space-level interaction and cross-scale feature pixel-level interaction on the multiple global dependent features to obtain target features, and determine a surface defect segmentation map of the rail surface image based on the target features;

[0040] The cross-scale feature channel-level interaction is to sequentially perform channel splitting, channel interaction fusion, and channel splicing on two adjacent global dependency features to obtain two first enhancement results;

[0041] The cross-scale feature space-level interaction is to normalize the size of the two first enhancement results, and then perform length splitting, spatial interaction fusion and length splicing in sequence to obtain two second enhancement results;

[0042] The cross-scale feature pixel-level interaction is to normalize the size of the two second enhancement results, and then perform feature interaction fusion based on the channel-level self-attention mechanism to obtain two third enhancement results, and all third enhancement results constitute the target feature.

[0043] The present invention also provides an electronic device, comprising a memory, a processor, and a computer program stored in the memory and executable on the processor, wherein when the processor executes the program, the rail surface defect segmentation method based on ordered cross-scale feature interaction as described in any one of the above-mentioned methods is implemented.

[0044] The present invention also provides a non-transitory computer-readable storage medium having a computer program stored thereon, which, when executed by a processor, implements any of the above-described methods for segmenting rail surface defects based on ordered cross-scale feature interaction.

[0045] The present invention also provides a computer program product, comprising a computer program, wherein when the computer program is executed by a processor, the method for segmenting rail surface defects based on ordered cross-scale feature interaction as described above is implemented.

[0046] Compared with the prior art, the present invention has the following beneficial effects:

[0047] The rail surface defect segmentation method and device based on ordered cross-scale feature interaction provided by the present invention can realize global dependent feature extraction, cross-scale feature channel-level interaction, feature space-level interaction and feature pixel-level interaction in sequence through the rail surface defect segmentation model, which can powerfully express the defects on the rail surface, fully enhance the features, highlight the defect areas in the features, and suppress irrelevant background areas. It is more suitable for defect features with complex backgrounds, low foreground-background contrast, and blurred target edges. The obtained features are more conducive to the subsequent defect segmentation process. Furthermore, by using the target features, a more accurate and detailed surface defect segmentation map of the rail surface image can be obtained, and accurate pixel-level positioning of the rail surface defects can be achieved. BRIEF DESCRIPTION OF THE DRAWINGS

[0048] In order to more clearly illustrate the technical solutions in the present invention or the prior art, the drawings required for use in the embodiments or the description of the prior art will be briefly introduced below. Obviously, for ordinary technicians in this field, other drawings can be obtained based on the drawings in the following description without paying any creative work.

[0049] Figure 1 It is a flow chart of a rail surface defect segmentation method based on ordered cross-scale feature interaction provided by the present invention;

[0050] Figure 2 It is a flow chart of the cross-scale feature channel-level interaction in the rail surface defect segmentation method based on ordered cross-scale feature interaction provided by the present invention;

[0051] Figure 3 It is a schematic flow chart of the spatial level interaction of cross-scale features in the rail surface defect segmentation method based on ordered cross-scale feature interaction provided by the present invention;

[0052] Figure 4 It is a flow chart of pixel-level interaction of cross-scale features in the rail surface defect segmentation method based on ordered cross-scale feature interaction provided by the present invention;

[0053] Figure 5 It is a structural schematic diagram of a rail surface defect segmentation device based on ordered cross-scale feature interaction provided by the present invention;

[0054] Figure 6 It is a structural schematic diagram of the electronic device provided by the present invention. DETAILED DESCRIPTION

[0055] In order to make the purpose, technical solution and advantages of the present invention clearer, the technical solution of the present invention will be clearly and completely described below in conjunction with the drawings of the present invention. Obviously, the described embodiments are part of the embodiments of the present invention, not all of the embodiments. Based on the embodiments of the present invention, all other embodiments obtained by ordinary technicians in this field without creative work are within the scope of protection of the present invention.

[0056] Figure 1 Schematic diagram of a flow chart of a rail surface defect segmentation method based on ordered cross-scale feature interaction provided in an embodiment of the present invention, such as Figure 1 As shown, the method includes:

[0057] S1, obtaining rail surface image;

[0058] S2, based on the rail surface defect segmentation model, extracting global dependent features from the rail surface image to obtain multiple global dependent features of different scales, and sequentially performing cross-scale feature channel-level interaction, cross-scale feature space-level interaction, and cross-scale feature pixel-level interaction on the multiple global dependent features to obtain target features, and determining a surface defect segmentation map of the rail surface image based on the target features;

[0059] The cross-scale feature channel-level interaction is to sequentially perform channel splitting, channel interaction fusion, and channel splicing on two adjacent global dependency features to obtain two first enhancement results;

[0060] The cross-scale feature space-level interaction is to normalize the size of the two first enhancement results, and then perform length splitting, spatial interaction fusion and length splicing in sequence to obtain two second enhancement results;

[0061] The cross-scale feature pixel-level interaction is to normalize the size of the two second enhancement results, and then perform feature interaction fusion based on the channel-level self-attention mechanism to obtain two third enhancement results, and all third enhancement results constitute the target feature.

[0062] Specifically, the rail surface defect segmentation method based on ordered cross-scale feature interaction provided in an embodiment of the present invention is executed by a rail surface defect segmentation device based on ordered cross-scale feature interaction, and the device can be configured in a computer, which can be a local computer or a cloud computer. The local computer can be a computer, a tablet, etc., which is not specifically limited here.

[0063] First, step S1 is performed to obtain a rail surface image. The rail surface image may be a rail surface image for which it is necessary to determine whether there are defects and the location of the defects. The rail surface image may be obtained by photographing the rail surface using an image acquisition device, which may be an industrial camera.

[0064] Then, step S2 is executed to introduce a rail surface defect segmentation model, and the surface defect information of the rail surface image is determined pixel by pixel through the rail surface segmentation model to obtain a surface defect segmentation map of the rail surface image.

[0065] The rail surface defect segmentation model may include a feature extraction network, a channel-level interaction module, a space-level interaction module, a pixel-level interaction module and a surface defect reasoning network connected in sequence.

[0066] The feature extraction network may be a PVT (Pyramid Vision Transformer) network, which is used to extract global dependent features from the rail surface image to obtain four initial features of different scales. The four initial features may be used as four global dependent features of different scales for extraction.

[0067] The channel-level interaction module can perform cross-scale feature channel-level interaction on multiple global dependency features, the spatial-level interaction module can perform cross-scale feature spatial-level interaction on the output of the channel-level interaction module, and the pixel-level interaction module can perform cross-scale feature pixel-level interaction on the spatial-level interaction module to obtain the target feature and use the target feature as output.

[0068] The surface defect reasoning network can use the target features to determine the surface defect segmentation map of the rail surface image.

[0069] Among them, the cross-scale feature channel-level interaction implemented by the channel-level interaction module includes: splitting two adjacent global dependency features into channels, fusing them by channel interaction, and splicing them into channels in turn to obtain two first enhancement results. Here, since the PVT network can obtain four global dependency features, each two adjacent global dependency features constitute a group, and a total of three groups can be obtained. Furthermore, the channel-level interaction module can obtain three groups of first enhancement results, each group of first enhancement results includes two first enhancement results.

[0070] The cross-scale feature spatial-level interaction implemented by the spatial-level interaction module includes: after normalizing the size of the two first enhancement results of each group of first enhancement results, length splitting, spatial interaction fusion and length splicing are performed in sequence to obtain two second enhancement results. Here, the spatial-level interaction module can obtain three groups of second enhancement results, each group of second enhancement results includes two second enhancement results.

[0071] The cross-scale feature pixel-level interaction implemented by the pixel-level interaction module includes: after normalizing the size of each group of two second enhancement results, performing feature interaction fusion based on the channel-level self-attention mechanism to obtain two third enhancement results, and all third enhancement results constitute the target feature. Here, the pixel-level interaction module can obtain three groups of third enhancement results, each group of third enhancement results includes two third enhancement results.

[0072] The rail surface defect segmentation method based on ordered cross-scale feature interaction provided in the embodiment of the present invention realizes global dependent feature extraction, cross-scale feature channel-level interaction, feature space-level interaction and feature pixel-level interaction in sequence through the rail surface defect segmentation model, which can powerfully express the defects on the rail surface, fully enhance the features, highlight the defect areas in the features, and suppress irrelevant background areas. It is more suitable for defect features with complex backgrounds, low foreground-background contrast, and blurred target edges. The obtained features are more conducive to the subsequent defect segmentation process. Furthermore, by using the target features, a more accurate and detailed surface defect segmentation map of the rail surface image can be obtained, and accurate pixel-level positioning of the rail surface defects can be achieved.

[0073] Based on the above embodiment, the cross-scale feature channel-level interaction specifically includes:

[0074] Perform channel splitting on the two global dependency features respectively to obtain two first sub-feature sets;

[0075] The first sub-features in the two first sub-feature sets are grouped accordingly to obtain multiple first sub-feature groups, and the two first sub-features in each first sub-feature group are subjected to spatial global average pooling and spatial global maximum pooling, and the global average pooling results and the global maximum pooling results of the two first sub-features are subjected to channel splicing, channel shuffling, two full connection operations, pixel-by-pixel summation, activation, and channel splitting in sequence to obtain first split features and second split features, and the first split features and the second split features are respectively multiplied pixel by pixel with the corresponding first sub-features to obtain two first enhanced features;

[0076] The two first enhancement features corresponding to each first sub-feature group are channel-joined in sequence to obtain the two first enhancement results.

[0077] Specifically, when performing cross-scale feature channel-level interaction, two adjacent global dependency features can be split into channels, each global dependency feature can be divided into four first sub-features with a channel number of 32, and the four first sub-features corresponding to each global dependency feature can constitute a first sub-feature set. Two global dependency features can correspond to two first sub-feature sets.

[0078] The first sub-features at corresponding positions in the two first sub-feature sets are grouped into one group to form four first sub-feature groups.

[0079] like Figure 2 As shown, the two first sub-features in each first sub-feature group can be represented as f 1 and f 2 , for f 1 and f 2 Both spatial-level global average pooling and spatial-level global maximum pooling are performed, and f 1 The global average pooling result is similar to f 2 The global average pooling result is sequentially subjected to channel splicing, channel shuffling, and two fully connected operations implemented by two fully connected layers to obtain the first result.

[0080] f 1 The global maximum pooling result and f 2 The global maximum pooling result is sequentially subjected to channel splicing, channel shuffling, and two fully connected operations implemented by two fully connected layers to obtain the second result.

[0081] The first result and the second result are summed pixel by pixel to obtain a third result, and the third result is activated by a sigmoid activation function and then channel-split to obtain a first split feature and a second split feature.

[0082] Combine the first split feature with f 1 Multiply pixel by pixel to get the same value as f 1 The corresponding first enhanced feature f 11 , the second split feature and f 2 Multiply pixel by pixel to get the same value as f 2 The corresponding first enhanced feature f 21 .

[0083] The above operation is performed on the two first sub-features in each first sub-feature group, and two first enhanced features corresponding to each first sub-feature group can be obtained. The two first enhanced features corresponding to each first sub-feature group are channel-joined in the order of the two first sub-feature sets to obtain two first enhanced results. Here, each global dependency feature corresponds to one first enhanced result.

[0084] In an embodiment of the present invention, through cross-scale feature channel-level interaction, two adjacent global dependent features can interact in the channel dimension, and feature enhancement is performed in the channel dimension, thereby obtaining a first enhancement result in the channel dimension, which can provide enhanced features in the channel dimension for subsequent determination of an accurate surface defect segmentation map.

[0085] Based on the above embodiment, the cross-scale feature space level interaction specifically includes:

[0086] After normalizing the sizes of the two first enhancement results respectively, splitting them in the length dimension to obtain two second sub-feature sets;

[0087] The second sub-features in the two second sub-feature sets are grouped accordingly to obtain multiple second sub-feature groups, and the two second sub-features in each second sub-feature group are spatially spliced ​​and then channel-level global average pooling and channel-level global maximum pooling are performed, and the global average pooling results and the global maximum pooling results of the two second sub-features are sequentially subjected to channel splicing, convolution operation, activation, and spatial splitting to obtain a third split feature and a fourth split feature, and the third split feature and the fourth split feature are respectively pixel-by-pixel multiplied with the corresponding second sub-features to obtain two second enhanced features;

[0088] The two second enhancement features corresponding to each second sub-feature group are respectively spliced ​​in length in sequence to obtain the two second enhancement results.

[0089] Specifically, when performing cross-scale feature space-level interaction, the two first enhancement results f obtained by cross-scale feature channel-level interaction can be 11 and f 21 After normalizing the sizes, the first enhancement results are split in the length dimension. Each scale-normalized first enhancement result can generate two second sub-features with a length half the length of the first enhancement feature. These two second sub-features can constitute a second sub-feature set. Therefore, the two first enhancement results f 11 and f 21 There may be two corresponding second sub-feature sets.

[0090] The second sub-features at corresponding positions in the two second sub-feature sets are grouped together to form two second sub-feature groups.

[0091] like Figure 3 As shown, the two second sub-features in each second sub-feature group can be represented as f 3 and f 4 , for f 3 and f 4 After spatial splicing, channel-level global average pooling and channel-level global maximum pooling are performed. The global average pooling results and global maximum pooling results of the two second sub-features are sequentially subjected to channel splicing, convolution operation implemented by a convolution layer, activation implemented by the sigmoid activation function, and spatial splitting to obtain the third split feature and the fourth split feature.

[0092] Combine the third split feature with f 3 Multiply pixel by pixel to get the same value as f 3 The corresponding second enhanced feature f 31, and the fourth split feature and f 4 Multiply pixel by pixel to get the same value as f 4 The corresponding second enhanced feature f 41 .

[0093] The above operation is performed on the two second sub-features in each second sub-feature group, and two second enhanced features corresponding to each second sub-feature group can be obtained. The two second enhanced features corresponding to each second sub-feature group are concatenated in length according to the order in the two second sub-feature sets to obtain two second enhanced results. Here, each first enhanced result corresponds to one second enhanced result.

[0094] In an embodiment of the present invention, through cross-scale feature space-level interaction, the two first enhancement results can interact in the spatial dimension, and feature enhancement is performed in the spatial dimension to obtain a second enhancement result in the spatial dimension, which can provide enhanced features in the spatial dimension for subsequent determination of an accurate surface defect segmentation map.

[0095] Based on the above embodiment, the cross-scale feature pixel-level interaction specifically includes:

[0096] Normalizing the sizes of the two second enhancement results respectively to obtain a third feature and a fourth feature;

[0097] After downsampling the third feature, two convolution operations are respectively performed to obtain a first value matrix and a query matrix, and two convolution operations are respectively performed to obtain a key matrix and a second value matrix on the fourth feature;

[0098] Transposing the query matrix to obtain a transposed matrix, and calculating an attention score based on the transposed matrix and the key matrix;

[0099] Transpose the attention score and perform matrix multiplication with the first value matrix after passing a normalized exponential function to obtain a fifth feature, and perform matrix multiplication with the second value matrix after passing a normalized exponential function to obtain a sixth feature;

[0100] Upsampling the fifth feature to obtain an upsampling result, and summing the upsampling result with the third feature to obtain a first summation result, and summing the sixth feature with the fourth feature to obtain a second summation result;

[0101] The first summation result and the second summation result are respectively subjected to convolution operations to obtain the two third enhancement results.

[0102] Specifically, when performing cross-scale feature pixel-level interaction, the two second enhancement results f obtained by cross-scale feature space-level interaction can be 31 and f 41Normalize the sizes respectively to get the third feature And the fourth characteristic .

[0103] For the third feature After downsampling, two convolution operations are performed to obtain the first value matrix V1 and the query matrix Q, and the fourth feature After two convolution operations, the key matrix K and the second value matrix V2 are obtained.

[0104] Transpose the query matrix Q to obtain the transposed matrix Q T , and the transposed matrix Q T Multiplying with the key matrix K gives the attention score A.

[0105] The attention score A is transposed and subjected to the softmax normalized exponential function, and then matrix multiplied with the first value matrix V1 to obtain the fifth feature.

[0106] The attention score A is normalized by the softmax exponential function and then matrix multiplied with the second value matrix V2 to obtain the sixth feature β.

[0107] The fifth feature is upsampled to obtain the upsampled result α, and the upsampled result α is added to the third feature Sum the result to get the first summation result, and add the sixth feature β to the fourth feature The sum is performed to obtain a second summation result.

[0108] Convolve the first summation result and the second summation result respectively to obtain the third feature The corresponding third enhancement result And the fourth characteristic The corresponding third enhancement result .

[0109] In an embodiment of the present invention, through cross-scale feature pixel-level interaction, the two second enhancement results can interact in the pixel dimension, and feature enhancement is performed in the pixel dimension to obtain a third enhancement result in the pixel dimension, which can provide pixel-dimensional enhancement features for subsequent determination of an accurate surface defect segmentation map.

[0110] On the basis of the above embodiment, determining the surface defect segmentation map of the rail surface image based on the target feature includes:

[0111] Based on the surface defect reasoning network in the rail surface defect segmentation model, the target feature is processed to obtain the surface defect segmentation map;

[0112] The surface defect reasoning network includes two convolutional layers, a random inactivation layer and a deconvolution layer connected in sequence.

[0113] Specifically, the rail surface defect segmentation model may include a surface defect reasoning network, through which the target features are processed to obtain a surface defect segmentation map.

[0114] The surface defect reasoning network may include two sequentially connected convolutional layers with a convolution kernel of 3×3, a random dropout layer with a hyperparameter of 0.5, and a deconvolution layer.

[0115] The surface defect reasoning network can be used to infer the surface defects of the target features and then obtain a surface defect segmentation map.

[0116] On the basis of the above-mentioned embodiment, the rail surface defect segmentation model is used to extract global dependent features from the rail surface image to obtain multiple global dependent features of different scales, including:

[0117] Based on the PVT network in the rail surface defect segmentation model, globally dependent features are extracted from the rail surface image to obtain four initial features of different scales;

[0118] The number of channels of the four initial features are unified to obtain four global dependency features.

[0119] Specifically, in the rail surface defect segmentation model, the feature extraction network can be a PVT network, and the four basic modules of Stage1, Stage2, Stage3 and Stage4 of the PVT network are used to extract global dependent features of the rail surface image to obtain four initial features of different scales. The introduction of the PVT network can obtain multiple features of different scales, which helps to obtain more comprehensive feature information.

[0120] Afterwards, the number of channels of the four initial features can be unified to 128 through a convolutional layer to obtain four global dependency features of different scales, so as to facilitate subsequent cross-scale feature interaction.

[0121] Based on the above embodiment, the rail surface defect segmentation model is trained based on the following steps:

[0122] Determine a training sample set, wherein the training sample set includes rail surface defect image samples and their corresponding segmentation truth images;

[0123] Inputting the rail surface defect image sample into an initial segmentation model to obtain a segmentation result output by the initial segmentation model;

[0124] Based on the segmentation result and the segmentation truth map, a binary cross entropy loss and an intersection-over-union loss are calculated, and based on the binary cross entropy loss and the intersection-over-union loss, a total loss is calculated;

[0125] Based on the total loss, the structural parameters of the initial segmentation model are iteratively updated using a gradient descent method to obtain the rail surface defect segmentation model.

[0126] Specifically, in the process of training the rail surface defect segmentation model, the training sample set is first determined. The training sample set can be constructed by taking 2086 artificial rail surface defect images in the NRSD-MN dataset as rail surface defect image samples and their corresponding segmentation truth maps.

[0127] When constructing the training sample set, the rail surface defect image samples and their corresponding segmentation truth maps can be rotated and flipped to perform data expansion, and their sizes can be adjusted to 3×352×352.

[0128] Thereafter, the rail surface defect image samples may be input into the initial segmentation model to obtain the segmentation results output by the initial segmentation model.

[0129] Using the segmentation results and the segmentation truth map, the binary cross entropy loss and the intersection-over-union loss are calculated, and the binary cross entropy loss and the intersection-over-union loss are weighted summed to obtain the total loss.

[0130] Using the total loss, the gradient descent method is used and the initial learning rate is set to 10 -4 The learning rate decay cycle is 1 / 10 of the original learning rate for every 30 complete traversals (Epoch) of the training sample set, and the batch size (Batchsize) is 16. The structural parameters of the initial segmentation model are iteratively updated until the preset number of iterations is reached or the total loss converges to obtain the rail surface defect segmentation model. Here, the preset number of iterations can be 70 Epochs.

[0131] Thereafter, a test sample set may also be determined, which may be constructed by using 965 artificial rail surface defect images and natural rail surface defect images in the NRSD-MN dataset as rail surface defect image samples through rail surface defect image samples and their corresponding segmentation truth maps.

[0132] During the test, the size of the rail surface defect image sample was adjusted to 3×352×352, and then input into the rail surface defect segmentation model to obtain the segmentation result. Then, the size of the segmentation result was adjusted to the same size as the input rail surface defect image sample.

[0133] Thereafter, the segmentation results obtained during the test process can be used to calculate the test index using the segmentation truth map corresponding to the rail surface defect image sample, and the test index can be used to determine the performance of the rail surface defect segmentation model.

[0134] The simulation experiment in the embodiment of the present invention is implemented based on PyTorch framework programming on a personal computer (PC) test platform with a central processing unit (CPU) of 3.0 GHz, a memory of 128 GB, a graphics processing unit (GPU) of RTX 3090, and a video memory of 24 GB.

[0135] like Figure 5 As shown, based on the above embodiment, an embodiment of the present invention provides a rail surface defect segmentation device based on ordered cross-scale feature interaction, comprising:

[0136] An image acquisition module 51 is used to acquire a rail surface image;

[0137] The defect segmentation module 52 is used to extract global dependent features of the rail surface image based on the rail surface defect segmentation model to obtain multiple global dependent features of different scales, and sequentially perform cross-scale feature channel-level interaction, cross-scale feature space-level interaction and cross-scale feature pixel-level interaction on the multiple global dependent features to obtain target features, and determine a surface defect segmentation map of the rail surface image based on the target features;

[0138] The cross-scale feature channel-level interaction is to sequentially perform channel splitting, channel interaction fusion, and channel splicing on two adjacent global dependency features to obtain two first enhancement results;

[0139] The cross-scale feature space-level interaction is to normalize the size of the two first enhancement results, and then perform length splitting, spatial interaction fusion and length splicing in sequence to obtain two second enhancement results;

[0140] The cross-scale feature pixel-level interaction is to normalize the size of the two second enhancement results, and then perform feature interaction fusion based on the channel-level self-attention mechanism to obtain two third enhancement results, and all third enhancement results constitute the target feature.

[0141] Based on the above embodiment, the defect segmentation module is specifically used for:

[0142] Perform channel splitting on the two global dependency features respectively to obtain two first sub-feature sets;

[0143] The first sub-features in the two first sub-feature sets are grouped accordingly to obtain multiple first sub-feature groups, and the two first sub-features in each first sub-feature group are subjected to spatial global average pooling and spatial global maximum pooling, and the global average pooling results and the global maximum pooling results of the two first sub-features are subjected to channel splicing, channel shuffling, two full connection operations, pixel-by-pixel summation, activation, and channel splitting in sequence to obtain first split features and second split features, and the first split features and the second split features are respectively multiplied pixel by pixel with the corresponding first sub-features to obtain two first enhanced features;

[0144] The two first enhancement features corresponding to each first sub-feature group are channel-joined in sequence to obtain the two first enhancement results.

[0145] Based on the above embodiment, the defect segmentation module is specifically used for:

[0146] After normalizing the sizes of the two first enhancement results respectively, splitting them in the length dimension to obtain two second sub-feature sets;

[0147] The second sub-features in the two second sub-feature sets are grouped accordingly to obtain multiple second sub-feature groups, and the two second sub-features in each second sub-feature group are spatially spliced ​​and then channel-level global average pooling and channel-level global maximum pooling are performed, and the global average pooling results and the global maximum pooling results of the two second sub-features are sequentially subjected to channel splicing, convolution operation, activation, and spatial splitting to obtain a third split feature and a fourth split feature, and the third split feature and the fourth split feature are respectively pixel-by-pixel multiplied with the corresponding second sub-features to obtain two second enhanced features;

[0148] The two second enhancement features corresponding to each second sub-feature group are respectively spliced ​​in length in sequence to obtain the two second enhancement results.

[0149] Based on the above embodiment, the defect segmentation module is specifically used for:

[0150] Normalizing the sizes of the two second enhancement results respectively to obtain a third feature and a fourth feature;

[0151] After downsampling the third feature, two convolution operations are respectively performed to obtain a first value matrix and a query matrix, and two convolution operations are respectively performed to obtain a key matrix and a second value matrix on the fourth feature;

[0152] Transposing the query matrix to obtain a transposed matrix, and calculating an attention score based on the transposed matrix and the key matrix;

[0153] Transpose the attention score and perform matrix multiplication with the first value matrix after passing a normalized exponential function to obtain a fifth feature, and perform matrix multiplication with the second value matrix after passing a normalized exponential function to obtain a sixth feature;

[0154] Upsampling the fifth feature to obtain an upsampling result, and summing the upsampling result with the third feature to obtain a first summation result, and summing the sixth feature with the fourth feature to obtain a second summation result;

[0155] The first summation result and the second summation result are respectively subjected to convolution operations to obtain the two third enhancement results.

[0156] Based on the above embodiment, the defect segmentation module is specifically used for:

[0157] Based on the surface defect reasoning network in the rail surface defect segmentation model, the target feature is processed to obtain the surface defect segmentation map;

[0158] The surface defect reasoning network includes two convolutional layers, a random inactivation layer and a deconvolution layer connected in sequence.

[0159] Based on the above embodiment, the defect segmentation module is specifically used for:

[0160] Based on the PVT network in the rail surface defect segmentation model, globally dependent features are extracted from the rail surface image to obtain four initial features of different scales;

[0161] The number of channels of the four initial features are unified to obtain four global dependency features.

[0162] On the basis of the above-mentioned embodiment, the rail surface defect segmentation device based on ordered cross-scale feature interaction further includes a model training module, which is specifically used for:

[0163] Determine a training sample set, wherein the training sample set includes rail surface defect image samples and their corresponding segmentation truth images;

[0164] Inputting the rail surface defect image sample into an initial segmentation model to obtain a segmentation result output by the initial segmentation model;

[0165] Based on the segmentation result and the segmentation truth map, a binary cross entropy loss and an intersection-over-union loss are calculated, and based on the binary cross entropy loss and the intersection-over-union loss, a total loss is calculated;

[0166] Based on the total loss, the structural parameters of the initial segmentation model are iteratively updated using a gradient descent method to obtain the rail surface defect segmentation model.

[0167] Specifically, the functions of each module in the rail surface defect segmentation device based on ordered cross-scale feature interaction provided in the embodiment of the present invention correspond one-to-one to the operating procedures of each step in the above-mentioned method embodiment, and the effects achieved are also consistent. Please refer to the above-mentioned embodiment for details, and no further details will be given in the embodiment of the present invention.

[0168] Figure 6 An example of a physical structure diagram of an electronic device is shown in FIG. Figure 6 As shown, the electronic device may include: a processor (Processor) 610, a communication interface (Communications Interface) 620, a memory (Memory) 630 and a communication bus 640, wherein the processor 610, the communication interface 620, and the memory 630 communicate with each other through the communication bus 640. The processor 610 may call the logic instructions in the memory 630 to execute the rail surface defect segmentation method based on the ordered cross-scale feature interaction provided in the above embodiments.

[0169] In addition, the logic instructions in the above-mentioned memory 630 can be implemented in the form of a software functional unit and can be stored in a computer-readable storage medium when it is sold or used as an independent product. Based on this understanding, the technical solution of the present invention is essentially or the part that contributes to the prior art or the part of the technical solution can be embodied in the form of a software product, and the computer software product is stored in a storage medium, including a number of instructions to enable a computer device (which can be a personal computer, a server, or a network device, etc.) to perform all or part of the steps of the method described in each embodiment of the present invention. The aforementioned storage medium includes: U disk, mobile hard disk, read-only memory (Read-Only Memory, ROM), random access memory (Random Access Memory, RAM), disk or optical disk and other media that can store program codes.

[0170] On the other hand, the present invention also provides a computer program product, which includes a computer program. The computer program can be stored on a non-transitory computer-readable storage medium. When the computer program is executed by a processor, the computer can execute the rail surface defect segmentation method based on ordered cross-scale feature interaction provided in the above embodiments.

[0171] On the other hand, the present invention also provides a non-transitory computer-readable storage medium having a computer program stored thereon, which, when executed by a processor, is implemented to execute the rail surface defect segmentation method based on ordered cross-scale feature interaction provided in the above-mentioned embodiments.

[0172] The device embodiments described above are merely illustrative, wherein the units described as separate components may or may not be physically separated, and the components displayed as units may or may not be physical units, that is, they may be located in one place, or they may be distributed on multiple network units. Some or all of the modules may be selected according to actual needs to achieve the purpose of the scheme of this embodiment. Ordinary technicians in this field can understand and implement it without paying creative labor.

[0173] Through the description of the above implementation methods, those skilled in the art can clearly understand that each implementation method can be implemented by means of software plus a necessary general hardware platform, and of course, can also be implemented by hardware. Based on this understanding, the above technical solution is essentially or the part that contributes to the prior art can be embodied in the form of a software product, and the computer software product can be stored in a computer-readable storage medium, such as ROM / RAM, a disk, an optical disk, etc., including a number of instructions for a computer device (which can be a personal computer, a server, or a network device, etc.) to execute the methods described in each embodiment or some parts of the embodiments.

[0174] Finally, it should be noted that the above embodiments are only used to illustrate the technical solutions of the present invention, rather than to limit it. Although the present invention has been described in detail with reference to the aforementioned embodiments, those skilled in the art should understand that they can still modify the technical solutions described in the aforementioned embodiments, or make equivalent replacements for some of the technical features therein. However, these modifications or replacements do not deviate the essence of the corresponding technical solutions from the spirit and scope of the technical solutions of the embodiments of the present invention.

Claims

1. A rail surface defect segmentation method based on ordered cross-scale feature interaction, characterized in that: include: Acquire rail surface images; Based on the rail surface defect segmentation model, globally dependent features are extracted from the rail surface image to obtain multiple global dependent features of different scales, and cross-scale feature channel-level interaction, cross-scale feature space-level interaction, and cross-scale feature pixel-level interaction are sequentially performed on the multiple global dependent features to obtain target features, and a surface defect segmentation map of the rail surface image is determined based on the target features; The cross-scale feature channel-level interaction is to sequentially perform channel splitting, channel interaction fusion, and channel splicing on two adjacent global dependency features to obtain two first enhancement results; The cross-scale feature space-level interaction is to normalize the size of the two first enhancement results, and then perform length splitting, spatial interaction fusion and length splicing in sequence to obtain two second enhancement results; The cross-scale feature pixel-level interaction is to normalize the size of the two second enhancement results, and then perform feature interaction fusion based on the channel-level self-attention mechanism to obtain two third enhancement results, and all the third enhancement results constitute the target feature; The cross-scale feature channel-level interaction specifically includes: Perform channel splitting on the two global dependency features respectively to obtain two first sub-feature sets; The first sub-features in the two first sub-feature sets are grouped accordingly to obtain multiple first sub-feature groups, and the two first sub-features in each first sub-feature group are subjected to spatial global average pooling and spatial global maximum pooling, and the global average pooling results and the global maximum pooling results of the two first sub-features are subjected to channel splicing, channel shuffling, two full connection operations, pixel-by-pixel summation, activation, and channel splitting in sequence to obtain first split features and second split features, and the first split features and the second split features are respectively multiplied pixel by pixel with the corresponding first sub-features to obtain two first enhanced features; Sequentially perform channel splicing on the two first enhancement features corresponding to each first sub-feature group to obtain the two first enhancement results; The cross-scale feature space-level interaction specifically includes: After normalizing the sizes of the two first enhancement results respectively, splitting them in the length dimension to obtain two second sub-feature sets; The second sub-features in the two second sub-feature sets are grouped accordingly to obtain multiple second sub-feature groups, and the two second sub-features in each second sub-feature group are spatially spliced ​​and then channel-level global average pooling and channel-level global maximum pooling are performed, and the global average pooling results and the global maximum pooling results of the two second sub-features are sequentially subjected to channel splicing, convolution operation, activation, and spatial splitting to obtain a third split feature and a fourth split feature, and the third split feature and the fourth split feature are respectively pixel-by-pixel multiplied with the corresponding second sub-features to obtain two second enhanced features; Sequentially concatenate the two second enhancement features corresponding to each second sub-feature group by length to obtain the two second enhancement results; The cross-scale feature pixel-level interaction specifically includes: Normalizing the sizes of the two second enhancement results respectively to obtain a third feature and a fourth feature; After downsampling the third feature, two convolution operations are respectively performed to obtain a first value matrix and a query matrix, and two convolution operations are respectively performed to obtain a key matrix and a second value matrix on the fourth feature; Transposing the query matrix to obtain a transposed matrix, and calculating an attention score based on the transposed matrix and the key matrix; Transpose the attention score and perform matrix multiplication with the first value matrix after passing a normalized exponential function to obtain a fifth feature, and perform matrix multiplication with the second value matrix after passing a normalized exponential function to obtain a sixth feature; Upsampling the fifth feature to obtain an upsampling result, and summing the upsampling result with the third feature to obtain a first summation result, and summing the sixth feature with the fourth feature to obtain a second summation result; The first summation result and the second summation result are respectively subjected to convolution operations to obtain the two third enhancement results.

2. The rail surface defect segmentation method based on ordered cross-scale feature interaction according to claim 1 is characterized in that: The step of determining a surface defect segmentation map of the rail surface image based on the target feature comprises: Based on the surface defect reasoning network in the rail surface defect segmentation model, the target feature is processed to obtain the surface defect segmentation map; The surface defect reasoning network includes two convolutional layers, a random inactivation layer and a deconvolution layer connected in sequence.

3. The rail surface defect segmentation method based on ordered cross-scale feature interaction according to claim 1 is characterized in that: Based on the rail surface defect segmentation model, the rail surface image is subjected to global dependency feature extraction to obtain multiple global dependency features of different scales, including: Based on the PVT network in the rail surface defect segmentation model, globally dependent features are extracted from the rail surface image to obtain four initial features of different scales; The number of channels of the four initial features are unified to obtain four global dependency features.

4. The rail surface defect segmentation method based on ordered cross-scale feature interaction according to any one of claims 1 to 3, characterized in that: The rail surface defect segmentation model is trained based on the following steps: Determine a training sample set, wherein the training sample set includes rail surface defect image samples and their corresponding segmentation truth images; Inputting the rail surface defect image sample into an initial segmentation model to obtain a segmentation result output by the initial segmentation model; Based on the segmentation result and the segmentation truth map, a binary cross entropy loss and an intersection-over-union loss are calculated, and based on the binary cross entropy loss and the intersection-over-union loss, a total loss is calculated; Based on the total loss, the structural parameters of the initial segmentation model are iteratively updated using a gradient descent method to obtain the rail surface defect segmentation model.

5. A rail surface defect segmentation device based on ordered cross-scale feature interaction, characterized in that: include: An image acquisition module, used for acquiring rail surface images; A defect segmentation module is used to extract global dependent features of the rail surface image based on the rail surface defect segmentation model to obtain multiple global dependent features of different scales, and sequentially perform cross-scale feature channel-level interaction, cross-scale feature space-level interaction and cross-scale feature pixel-level interaction on the multiple global dependent features to obtain target features, and determine a surface defect segmentation map of the rail surface image based on the target features; The cross-scale feature channel-level interaction is to sequentially perform channel splitting, channel interaction fusion, and channel splicing on two adjacent global dependency features to obtain two first enhancement results; The cross-scale feature space-level interaction is to normalize the size of the two first enhancement results, and then perform length splitting, spatial interaction fusion and length splicing in sequence to obtain two second enhancement results; The cross-scale feature pixel-level interaction is to normalize the size of the two second enhancement results, and then perform feature interaction fusion based on the channel-level self-attention mechanism to obtain two third enhancement results, and all the third enhancement results constitute the target feature; The cross-scale feature channel-level interaction specifically includes: Perform channel splitting on the two global dependency features respectively to obtain two first sub-feature sets; The first sub-features in the two first sub-feature sets are grouped accordingly to obtain multiple first sub-feature groups, and the two first sub-features in each first sub-feature group are subjected to spatial global average pooling and spatial global maximum pooling, and the global average pooling results and the global maximum pooling results of the two first sub-features are subjected to channel splicing, channel shuffling, two full connection operations, pixel-by-pixel summation, activation, and channel splitting in sequence to obtain first split features and second split features, and the first split features and the second split features are respectively multiplied pixel by pixel with the corresponding first sub-features to obtain two first enhanced features; Sequentially perform channel splicing on the two first enhancement features corresponding to each first sub-feature group to obtain the two first enhancement results; The cross-scale feature space-level interaction specifically includes: After normalizing the sizes of the two first enhancement results respectively, splitting them in the length dimension to obtain two second sub-feature sets; The second sub-features in the two second sub-feature sets are grouped accordingly to obtain multiple second sub-feature groups, and the two second sub-features in each second sub-feature group are spatially spliced ​​and then channel-level global average pooling and channel-level global maximum pooling are performed, and the global average pooling results and the global maximum pooling results of the two second sub-features are sequentially subjected to channel splicing, convolution operation, activation, and spatial splitting to obtain a third split feature and a fourth split feature, and the third split feature and the fourth split feature are respectively pixel-by-pixel multiplied with the corresponding second sub-features to obtain two second enhanced features; Sequentially concatenate the two second enhancement features corresponding to each second sub-feature group by length to obtain the two second enhancement results; The cross-scale feature pixel-level interaction specifically includes: Normalizing the sizes of the two second enhancement results respectively to obtain a third feature and a fourth feature; After downsampling the third feature, two convolution operations are respectively performed to obtain a first value matrix and a query matrix, and two convolution operations are respectively performed to obtain a key matrix and a second value matrix on the fourth feature; Transposing the query matrix to obtain a transposed matrix, and calculating an attention score based on the transposed matrix and the key matrix; Transpose the attention score and perform matrix multiplication with the first value matrix after passing a normalized exponential function to obtain a fifth feature, and perform matrix multiplication with the second value matrix after passing a normalized exponential function to obtain a sixth feature; Upsampling the fifth feature to obtain an upsampling result, and summing the upsampling result with the third feature to obtain a first summation result, and summing the sixth feature with the fourth feature to obtain a second summation result; The first summation result and the second summation result are respectively subjected to convolution operations to obtain the two third enhancement results.

6. 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 program, the rail surface defect segmentation method based on ordered cross-scale feature interaction as described in any one of claims 1 to 4 is implemented.

7. A non-transitory computer-readable storage medium having a computer program stored thereon, characterized in that: When the computer program is executed by a processor, the rail surface defect segmentation method based on ordered cross-scale feature interaction as described in any one of claims 1 to 4 is implemented.

Images