Metal surface defect image enhancement method, system and equipment based on residual multi-head attention

Abstract

The invention discloses a metal surface defect image enhancement method, system and equipment based on residual multi-head attention, and solves the problem that machine vision is not high in metal product surface tiny defect recognition precision, and the method comprises the following steps: S1, obtaining an initial defect image containing defect features; s2, inputting the initial defect image into a pre-trained residual multi-head attention fusion model, performing defect feature extraction on the input initial defect image by the residual multi-head attention fusion model, obtaining a feature sequence from an extracted defect feature map through a differential multi-head attention mechanism, and converting the feature sequence into an enhanced feature map of defect features; and S3, amplifying the enhanced feature map through sub-pixel convolution, and outputting defect image data after enhancement and amplification. According to the method, the defect features in the defect image are enhanced and losslessly amplified through the pre-trained residual multi-head attention fusion model, the small-size defect features in the image are more truly obtained, and the detection effect of the metal surface tiny defects is improved.

Description

Image enhancement method, system, and device for metal surface defects based on residual multi-head attention. Technical Field

[0001] This invention relates to a method, system, and device for enhancing images of metal surface defects based on residual multi-head attention, belonging to the interdisciplinary technologies of computer vision and machine learning. Background Technology

[0002] Metal products often exhibit surface defects such as scratches, corrosion, cracks, and dirt due to factors including raw materials, manufacturing processes, and casting environment. Many of these are minute defects that are difficult to detect. These minute defects can significantly reduce the performance and lifespan of metal products. Failure to accurately detect potential minute defects can lead to substantial economic losses and even incalculable safety accidents. Therefore, detecting minute defects on the surface of metal products is a crucial step in ensuring product quality.

[0003] Traditional non-destructive testing (NDT) methods include eddy current testing, magnetic particle testing, radiographic testing, ultrasonic testing, isotope photography, penetrant testing, and infrared thermography. These techniques utilize auxiliary tools to detect defects, and most require specialized testing equipment, placing high demands on the operators. Furthermore, while these traditional techniques are effective for routine defect detection, they still have a significant probability of missing subtle, minute defects, making them unsuitable for detecting such defects.

[0004] With the development of computer image processing technology, machine vision has become a simple and efficient method for detecting surface defects in metal products without the need for professional operation. However, for tiny defects within 1 cm, the defect feature size is too small and may only occupy a few pixels in the entire product surface image after imaging, which is not conducive to the subsequent identification of the defect location. Therefore, how to improve the detection efficiency and accuracy of tiny defects has become an urgent problem to be solved by machine vision in the detection of surface defects in metal products. Summary of the Invention

[0005] The technical problem solved by this invention is to address the issue of low accuracy in identifying minute defects on the surface of metal products using existing machine vision systems, and to provide a method, system, and device for enhancing metal surface defect images based on residual multi-head attention.

[0006] This invention is achieved using the following technical solution:

[0007] This invention first discloses a method for enhancing metal surface defect images based on residual multi-head attention, comprising the following steps:

[0008] S1. Obtain an initial defect image containing defect features;

[0009] S2. The initial defect image is input into a pre-trained residual multi-head attention fusion model. The residual multi-head attention fusion model extracts defect features from the input initial defect image, and obtains a feature sequence from the extracted defect feature map through a differential multi-head attention mechanism, which is then converted into an enhanced feature map of the defect features. The process of the multi-head attention mechanism is as follows:

[0010] The defect feature map is flattened into a one-dimensional vector. Two sets of linear transformations are used to generate two sets of queries, keys, and values ​​for the defect feature map. The two sets of queries, keys, and values ​​are then segmented. Attention scores for the segmented queries, keys, and values ​​are calculated using each attention head. The outputs of each attention head are then concatenated to convert the defect feature map into an enhanced feature map.

[0011] S3. The enhanced feature map is amplified through sub-pixel convolution, and the enhanced and amplified defect image data is output.

[0012] In the metal surface defect image enhancement method based on residual multi-head attention of the present invention, the residual module of the residual multi-head attention fusion model further adopts a Conv-BN-ReLU chain structure.

[0013] In the metal surface defect image enhancement method based on residual multi-head attention of the present invention, the attention score of the attention head in step S2 is further calculated by the following formula:

[0014] ,

[0015] ,

[0016] ,

[0017] Where A1 and A2 are two independent projections generated from the input feature sequence, d k For each attention head dimension, Q1' and Q2' are the two sets of queries after splitting, K1' and K2' are the two sets of keys after splitting, V1' and V2' are the two sets of values ​​after splitting, T is the matrix transpose, and head i This indicates the final output of the current header. This represents a learnable differential weight parameter.

[0018] In the metal surface defect image enhancement method based on residual multi-head attention of the present invention, step S3 further includes the following sub-steps:

[0019] S31. Based on the set magnification factor r value, use a standard convolution with a kernel size of k×k to generate a feature tensor with the number of channels equal to the square of the target magnification factor r value from the enhanced feature map of the defect features.

[0020] S32. A periodic filtering operation is used to map the deep-dimensional channel information of the defect features to the spatial dimension. The resulting feature tensor is then split according to the channel dimension to obtain r. 2 Each sub-tensor is rearranged according to a chessboard pattern to obtain lossless magnified defect image data, specifically represented as follows:

[0021] ,

[0022] Among them, Input (H,W,C) The enhanced feature map represents the input defect features, where C is the number of channels, H×W is the spatial dimension, and the number of channels is expanded to C×r. 2 Output (rH,rW,C) This represents the feature map of the defect image output, which is the enhanced and magnified defect image data output in this step.

[0023] In the metal surface defect image enhancement method based on residual multi-head attention of the present invention, the training process of the residual multi-head attention fusion model is further as follows:

[0024] K1. Obtain the defect image dataset, wherein the defect image samples in the defect image dataset contain specific defect feature information and defect feature annotation text information;

[0025] K2. Input the defect image samples in the defect image dataset into the residual multi-head attention fusion model. The residual multi-head attention fusion model extracts the sample defect features from the input defect image samples. The extracted sample defect feature map is processed by the differential multi-head attention mechanism to obtain the sample feature sequence and convert it into an enhanced feature map of the sample defect features.

[0026] K3. The enhanced feature map of the sample defect features is magnified through sub-pixel convolution, and the enhanced and magnified defect reconstruction image is output.

[0027] K4. The perceptual loss function is used to calculate the difference in feature maps between the reconstructed defect image and the real defect image in the defect image sample. Combined with adversarial loss, a feedback mechanism is formed to adjust the residual multi-head attention fusion model.

[0028] K5. The residual multi-head attention fusion model is evaluated and tested by using the defect feature annotation text information in the defect image samples.

[0029] In the metal surface defect image enhancement method based on residual multi-head attention of the present invention, further, in step K4, the perceptual loss function includes adversarial loss and content loss.

[0030] ,

[0031] in, For the perceptual loss function, For content loss function, To counteract the loss function, Weighting coefficients to counteract losses;

[0032] The adversarial loss function minimizes the difference between the defect-reconstructed image and the defect-real image in the discriminator output, prompting the generator to produce a realistic image, as shown below:

[0033] ,

[0034] Where x is the real high-resolution image, corresponding to the defect reconstruction image, L is the input low-resolution image, corresponding to the real defect image in the sample, G is the generator, G(L) is the defect image generated by the generator based on the input low-resolution image, D is the discriminator, D(x) determines the probability that x is the real image, D(G(L)) determines the probability that it is the generated image, and E[ ] represents the expectation;

[0035] The content loss function utilizes a feature extraction network to extract feature maps from the reconstructed defect image and the actual defect image, respectively, and calculates the difference between the two, expressed as:

[0036] ,

[0037] Wherein, FVGG(x) is the feature map of the defect-reconstructed image output by the feature extraction network. The feature map is the output of the feature extraction network to show the real image of the defect. W, H, and C are the width, height, and number of channels of the extracted feature map, respectively, and i, j, and k represent the height index, width index, and channel index within the feature map, respectively.

[0038] In the metal surface defect image enhancement method based on residual multi-head attention of the present invention, the defect image dataset is further derived from the NEU-DET dataset.

[0039] In the metal surface defect image enhancement method based on residual multi-head attention of the present invention, the feature extraction network is a pre-trained VGG-19 network model.

[0040] The present invention also discloses an image enhancement system based on residual multi-head attention that implements the above steps, comprising:

[0041] The image acquisition module acquires an initial defect image containing defect features;

[0042] The residual multi-head attention fusion module inputs the initial defect image into a pre-trained residual multi-head attention fusion model. The residual multi-head attention fusion model extracts defect features from the input initial defect image, and then uses a differential multi-head attention mechanism to obtain a feature sequence from the extracted defect feature map, converting it into an enhanced feature map of the defect features. The process of the multi-head attention mechanism is as follows:

[0043] The defect feature map is flattened into a one-dimensional vector. Two sets of linear transformations are used to generate two sets of queries, keys, and values ​​for the defect feature map. The two sets of queries, keys, and values ​​are then segmented. Attention scores for the segmented queries, keys, and values ​​are calculated using each attention head. The outputs of each attention head are then concatenated to convert the defect feature map into an enhanced feature map.

[0044] The subpixel convolution module amplifies the enhanced feature map through subpixel convolution and outputs the enhanced and amplified defect image data.

[0045] The present invention also discloses a computer device, including a memory and a processor, wherein the memory is used to store a computer program, and the processor runs the computer program to enable the electronic device to perform the above-described method for enhancing metal surface defect images based on residual multi-head attention.

[0046] The present invention, by adopting the above technical solution, has the following beneficial effects:

[0047] (1) This invention provides a residual multi-head attention fusion model to enhance the recognition of defect features in defect images. The residual multi-head attention fusion model includes a standard residual module and a multi-head attention module. The standard residual module retains the dual convolution structure of the standard residual module. It captures local texture features in the defect image through convolution, reshapes the defect feature map into a sequence form, and generates a query Q, key K, and value V matrix according to the linear layer with shared parameters. The subspace is split according to the number of heads in the multi-head attention module, and the attention score is calculated for each head. The outputs of multiple attention heads are concatenated and converted into an enhanced feature map of the defect features. This solves the problem that the local receptive field of traditional convolution is difficult to capture cross-regional correlation features. By combining the residual mechanism and the multi-head attention mechanism, long-range dependencies between pixels can be modeled through global interaction, and more defect feature information can be learned to enhance the quality of the reconstructed image. Sub-pixel convolution is used to perform lossless amplification processing on the defect feature enhancement obtained by the residual multi-head attention fusion model to obtain defect image data after setting the amplification factor, so as to support the network model in the task of recognizing and detecting small-sized defects.

[0048] (2) In the training process of the residual multi-head attention fusion model, the present invention uses the residual multi-head attention fusion model to focus on and learn the feature texture details of each sample in the defect image dataset. The differential multi-head attention module is used to suppress irrelevant noise in the defect feature map of each sample in the defect image dataset and further focus on learning the small defect feature information in the dataset samples to obtain the enhanced feature map of the enhanced defect image dataset samples. After lossless amplification of the enhanced feature map by setting the amplification factor using sub-pixel convolution, the perceptual loss function is used to solve the problem of loss of defect texture details caused by the smoothing phenomenon in the image reconstruction process. The perceptual loss function includes two parts: adversarial loss and content loss. The adversarial loss is based on generative adversarial game, which can enable the generator to generate realistic images and focus more on whether the generated image is close to the distribution of the real image. The content loss focuses on the differences in details and high-level semantic features between the generated image and the real image. The feature extraction model is used to extract the real defect features in the initial defect image and the generated defect features in the defect reconstruction image respectively. After calculating the difference between the two, the adversarial loss is combined to form a feedback mechanism to fine-tune the residual multi-head attention fusion model, so that the model obtains defect features with good reconstruction effect in the process of image enhancement.

[0049] In summary, the metal surface defect image enhancement method, system, and device provided by this invention enhances and non-destructively amplifies defect features in defect images through a pre-trained residual multi-head attention fusion model. This enables more realistic acquisition of small-sized defect features in images, improves defect detection performance, and provides support for defect recognition in images. It is particularly suitable for enhancing and recognizing tiny defects on metal surfaces, effectively reducing detection costs for detecting tiny defects on metal surfaces.

[0050] The present invention will be further described below with reference to the accompanying drawings and specific embodiments. Attached Figure Description

[0051] Figure 1 is a schematic diagram of the process of the metal surface defect image enhancement method based on residual multi-head attention according to the present invention.

[0052] Figure 2 is a schematic diagram of the training process of the residual multi-head attention fusion model in this invention.

[0053] Figure 3 is a schematic diagram of the metal surface defect image enhancement method based on residual multi-head attention according to the present invention.

[0054] Figure 4 is a schematic diagram of the principle of the residual multi-head attention fusion model in this invention to generate enhanced feature maps of defect features.

[0055] Figure 5 is a schematic diagram of the process of amplifying the enhanced feature map through sub-pixel convolution according to the present invention.

[0056] Figure 6 is a schematic diagram of the image enhancement system based on residual multi-head attention according to the present invention.

[0057] Figure 7 is a schematic diagram of a computer device for implementing the metal surface defect image enhancement method based on residual multi-head attention of the present invention. Detailed Implementation

[0058] The technical solutions of the embodiments of the present invention will be clearly and completely described below with reference to the accompanying drawings. Obviously, the described embodiments are only some embodiments of the present invention, and not all embodiments. Based on the embodiments of the present invention, all other embodiments obtained by those skilled in the art without creative effort are within the scope of protection of the present invention.

[0059] To make the above-mentioned objects, features and advantages of the present invention more apparent and understandable, the present invention will be further described in detail below with reference to the accompanying drawings and specific embodiments.

[0060] Referring to Figures 1 and 3, the illustrations show a flowchart of a metal surface defect image enhancement method based on residual multi-head attention according to the present invention, which specifically includes the following steps:

[0061] S1. Obtain an initial defect image containing defect features. Taking metal surface defect recognition as an example, obtain an initial defect image containing metal surface defect features.

[0062] S2. Input the initial defect image into the pre-trained residual multi-head attention fusion model, as shown in Figure 3. The residual multi-head attention fusion model is an integrated module including a standard residual module and a differential multi-head attention module. The residual module of the residual multi-head attention fusion model adopts a Conv-BN-ReLU chain structure. It extracts the metal surface defect features from the input initial defect image, and obtains the feature sequence from the extracted defect feature map through the differential multi-head attention mechanism and converts it into an enhanced feature map of metal surface defect features.

[0063] As shown in Figure 4, the process of the multi-head attention mechanism is as follows:

[0064] The defect feature map is flattened into a one-dimensional vector. Two sets of linear transformations are used to generate two sets of queries, keys, and values ​​for the defect feature map. These two sets of queries, keys, and values ​​are then segmented. The query, key, and value matrix is ​​divided into multiple heads along the feature dimension (channel dimension), each head having a smaller dimension. For example, if the original query has dimensions B × N × C, where N is the sequence length (H*W) and C is the number of channels, in multi-head attention, C is typically divided into h heads, each head having a dimension of d.k = C / h, where h is a user-defined value, set to 4 here. Attention scores are calculated for each segmented query, key, and value using each attention head. The outputs of each attention head are then concatenated to create an enhanced feature map of the defect features.

[0065] The attention score of the attention head is calculated using the following formula:

[0066] ,

[0067] ,

[0068] ,

[0069] Where A1 and A2 are two independent projections generated from the input feature sequence, d k For each attention head dimension, Q1' and Q2' are the two sets of queries after splitting, K1' and K2' are the two sets of keys after splitting, V1' and V2' are the two sets of values ​​after splitting, T is the matrix transpose, and head i This indicates the final output of the current header. This represents a learnable differential weight parameter.

[0070] S3. The enhanced feature map is amplified through sub-pixel convolution, and the enhanced and amplified defect image data is output.

[0071] The sub-pixel convolution is divided into a feature extraction stage and a pixel recombination stage, which outputs the enhanced feature map of the metal surface defect to obtain the metal surface defect image data after setting the magnification factor.

[0072] This step specifically includes the following sub-steps:

[0073] S31. Based on the set magnification factor r value, use a standard convolution with a kernel size of k×k to generate a feature tensor with the number of channels equal to the square of the target magnification factor r value from the enhanced feature map of the defect features.

[0074] S32. A periodic filtering operation is used to map the deep-dimensional channel information of the defect features to the spatial dimension. The resulting feature tensor is then split according to the channel dimension to obtain r. 2 Each sub-tensor is rearranged according to a chessboard pattern to obtain lossless magnified defect image data, specifically represented as follows:

[0075] ,

[0076] Among them, Input (H,W,C)The enhanced feature map represents the input defect features, where C is the number of channels and H×W is the spatial dimension. PixelShuffle() is the core operation of sub-pixel convolution, which rearranges the channel dimension information into the spatial dimension, thereby achieving image upsampling and magnification, expanding the number of channels to C×r. 2 Output (rH,rW,C) This represents the feature map of the defect image output, which is the enhanced and magnified defect image data output in this step.

[0077] Referring to Figure 2, the training process of the residual multi-head attention fusion model in this invention is as follows:

[0078] K1. Obtain the defect image dataset. The defect image samples in the dataset contain specific defect feature information and defect feature annotation text information. The defect image dataset is selected from open-source datasets. Taking the processing of metal surface defect images as an example, the NEU-DET dataset can be selected.

[0079] After downloading and obtaining the defect image dataset from the open-source dataset, the quality of the defect image samples in the obtained metal surface defect image dataset is evaluated, and poorly performing metal surface defect images are removed. Furthermore, based on the specific defect information and defect annotation text information contained in each metal surface defect image sample, the samples in the dataset are classified, and the image defect information and text annotation information are adjusted to ensure that the two information are accurate and consistent.

[0080] K2. A residual multi-head attention fusion model is used to focus on and learn the image information of metal surface defects. Defect image samples from the defect image dataset are input into the residual multi-head attention fusion model. The residual multi-head attention fusion model extracts sample defect features from the input defect image samples. The extracted sample defect feature maps are then processed through a differential multi-head attention mechanism to obtain sample feature sequences and convert them into enhanced feature maps of sample defect features.

[0081] The residual multi-head attention fusion model in this invention uses a chain-structured residual module of "Conv-BN-ReLU" to perform data processing and feature extraction on each defect image sample in the metal surface defect image dataset, and obtains the sample defect feature map corresponding to the defect image sample.

[0082] The acquired sample defect feature map is flattened into a one-dimensional vector. Where B is the batch size for training the defect image dataset, C is the number of channels, and H×W is the spatial dimension. Differential multi-head attention is used to suppress irrelevant noise and capture global dependencies and contextual information in the feature maps. Two sets of linear transformations are used to generate a query, key, and value, resulting in the enhanced feature sequence of the defect image dataset, which is then transformed into an enhanced feature map of the sample defect features. The specific process is as follows:

[0083] ;

[0084] .

[0085] Where Q1, K1, and V1 are the first group of query, key, and value, respectively, and W Q1 W K1 W V1 These are the learning weight matrices for the first group, Q2, K2, and V2, respectively, representing the query, key, and value for the second group. Q2 W K2 W V2 These are the learning weight matrices corresponding to the second group. Attention scores for the segmented query, key, and value are calculated using each attention head, and the outputs of each attention head are concatenated to convert them into an enhanced feature map of the defect features.

[0086] The attention score of the attention head is calculated using the following formula:

[0087] ,

[0088] ,

[0089] .

[0090] Where A1 and A2 are two independent projections generated from the input feature sequence, d k For each attention head dimension, Q1' and Q2' are the two sets of queries after splitting, K1' and K2' are the two sets of keys after splitting, V1' and V2' are the two sets of values ​​after splitting, T is the matrix transpose, and head i This indicates the final output of the current header. This represents a learnable differential weight parameter.

[0091] The entire process described above can be summarized as follows: . represents the output of multi-head attention on input feature X. Concat concatenates the outputs of all attention heads along the feature dimension. Reshape, representing the output of all attention heads, adjusts the tensor shape to ensure correct dimensions.

[0092] K3. The enhanced feature map of the sample defect features is amplified through sub-pixel convolution, and the enhanced and amplified defect reconstruction image is output. Convolution is used to learn the low-resolution to high-resolution feature mapping of the defect feature enhancement feature map, and then the information is reorganized from the channel dimension to the spatial dimension, as shown in Figure 5.

[0093] Based on the set magnification factor r, a standard convolution with a kernel size of k×k is used to generate a feature tensor with the number of channels equal to the square of the target magnification factor r from the enhanced feature map of the defect features.

[0094] Periodic filtering is used to map the deep-dimensional channel information of defect features to a spatial dimension, and the resulting feature tensor is decomposed according to the channel dimension to obtain r. 2 Each sub-tensor is rearranged according to a chessboard pattern to obtain lossless magnified defect image data. Specifically, it is represented as follows:

[0095] ,

[0096] Among them, Input (H,W,C) The enhanced feature map represents the input defect features, where C is the number of channels and H×W is the spatial dimension. PixelShuffle() is the core operation of sub-pixel convolution, which rearranges the channel dimension information into the spatial dimension, thereby achieving image upsampling and magnification, expanding the number of channels to C×r. 2 Output (rH,rW,C) This represents the feature map of the output defect image, which is the reconstructed defect image after enhancing and magnifying the defect image samples in the sample defect image dataset.

[0097] A periodic filtering operation is used to map deep-dimensional channel information to a spatial dimension. The resulting r-squared feature tensor is then split into r² sub-tensors according to the channel dimension, and rearranged in a chessboard pattern to obtain a lossless amplified dataset sample. The specific arrangement process is represented as follows:

[0098] ,

[0099] Where T is the input tensor (H, W, C×r2), x and y are the spatial coordinates of the output feature map, and the modulo operation is used to locate the specific position of the sub-pixel within the block.

[0100] K4. A perceptual loss function is used to calculate the difference in feature maps between the reconstructed defect image and the real defect image in the defect image sample. This difference is combined with adversarial loss to form a feedback mechanism to adjust the residual multi-head attention fusion model. The difference between the feature maps of the reconstructed image and the real image is calculated and combined with the adversarial loss of the network model. Based on this difference, the model is fine-tuned. The adversarial loss function is used to minimize the difference between the generated image and the real high-resolution image in the discriminator output, so as to achieve the generator producing realistic images.

[0101] The perceptual loss function consists of two parts: adversarial loss and content loss, as follows:

[0102] .

[0103] in, For the perceptual loss function, For content loss function, To counteract the loss function, The weighting factor to counteract the loss is set to 10. -3 This is used to balance the contribution of content loss and combat loss.

[0104] The adversarial loss function minimizes the difference between the defect-reconstructed image and the defect-real image in the discriminator output, prompting the generator to produce a realistic image, as shown below:

[0105] .

[0106] Where x is the real high-resolution image, corresponding to the defect reconstruction image, L is the input low-resolution image, corresponding to the real defect image in the sample, G is the generator, G(L) is the defect image generated by the generator based on the input low-resolution image, D is the discriminator, D(x) determines the probability that x is the real image, D(G(L)) determines the probability that it is the generated image, and E[ ] represents the expectation.

[0107] The content loss function utilizes a feature extraction network to extract feature maps from the reconstructed defect image and the real defect image, and calculates the difference between them. This greatly avoids the smoothing problem that exists in the model-generated reconstructed defect image, and achieves a better recovery effect on the model's feature detail texture information of the defect.

[0108] Taking the pre-trained VGG-19 network model as the feature extraction network as an example, the content loss function is expressed as:

[0109] .

[0110] Wherein, FVGG(x) is the feature map of the defect-reconstructed image output by the feature extraction network. The feature map is the output of the feature extraction network to show the real image of the defect. W, H, and C are the width, height, and number of channels of the extracted feature map, respectively. i, j, and k represent the height index, width index, and channel index within the feature map, respectively. Here, i is the height index (0 to H-1), j is the width index (0 to W-1), and k is the channel index (0 to C-1).

[0111] K5. Obtain the detection results of the pre-enhancement and post-enhancement images. Use the defect feature annotation text information in the defect image samples to evaluate and test the residual multi-head attention fusion model. Evaluate the model using two objective evaluation metrics: PSNR and SSIM. Test the model using defect image samples that have never participated in training or evaluation, as well as using subsequently generated images.

[0112] As shown in Figure 6, the image enhancement system based on residual multi-head attention of the present invention includes an image acquisition module, a residual multi-head attention fusion module, and a sub-pixel convolution module.

[0113] The image acquisition module acquires an initial defect image containing defect features.

[0114] The residual multi-head attention fusion module inputs the initial defect image into a pre-trained residual multi-head attention fusion model. The residual multi-head attention fusion model extracts defect features from the input initial defect image, and obtains a feature sequence from the extracted defect feature map through a differential multi-head attention mechanism, which is then converted into an enhanced feature map of the defect features. The process of the multi-head attention mechanism is as follows: the defect feature map is flattened into a one-dimensional vector, and two sets of queries, keys, and values ​​are generated from the defect feature map through two sets of linear transformations. The two sets of queries, keys, and values ​​are segmented, and the attention score of the segmented queries, keys, and values ​​is calculated by each attention head. The outputs of each attention head are concatenated to convert the model into an enhanced feature map of the defect features.

[0115] The residual multi-head attention fusion model directly passes the initial defect image to subsequent layers through skip connections. It utilizes differential multi-head attention modules to suppress irrelevant noise and focus on minor defects in the image. Finally, this output is added to the input to form a residual learning mechanism. This design mechanism solves the gradient vanishing and network degradation problems in deep network training, enabling the residual multi-head attention fusion model to more efficiently transmit gradient information from defect features while preserving input features, greatly aiding in defect feature extraction. The residual multi-head attention fusion model retains the double convolutional structure of the standard residual module, capturing local texture features of defects through convolution. The feature map is reshaped into a sequence, and a query Q, key K, and value V matrix is ​​generated based on a linear layer with shared parameters. The model is then subspaced according to the number of heads, and an attention score is calculated for each head. The outputs of multiple heads are concatenated and converted into an image format.

[0116] The specific process is as follows:

[0117] The acquired defect feature map is flattened into a one-dimensional vector. Where B is the batch size for training the defect image dataset, C is the number of channels, and H×W is the spatial dimension. A differential multi-head attention mechanism is used to suppress irrelevant noise in the feature maps and capture global dependencies and contextual information. Two sets of linear transformations are used to generate a query, key, and value, resulting in the feature sequence of the enhanced defect image dataset, which is then transformed into an enhanced feature map of the defect features. The specific process is as follows:

[0118] ;

[0119] .

[0120] Where Q1, K1, and V1 are the first group of query, key, and value, respectively, and W Q1 W K1 W V1 These are the learning weight matrices for the first group, Q2, K2, and V2, respectively, representing the query, key, and value for the second group. Q2 W K2 W V2 These are the learning weight matrices corresponding to the second group. Attention scores for the segmented query, key, and value are calculated using each attention head, and the outputs of each attention head are concatenated to convert them into an enhanced feature map of the defect features.

[0121] The attention score of the attention head is calculated using the following formula:

[0122] ,

[0123] ,

[0124] ,

[0125] Here, dk is the dimension of each attention head, Q1' and Q2' are the two sets of queries after segmentation, K1' and K2' are the two sets of keys after segmentation, and V1' and V2' are the two sets of values ​​after segmentation. Softmax is used for normalization, and finally, the outputs of each attention head are concatenated and converted into a feature map shape.

[0126] The entire process described above can be summarized as follows: This formula includes the specific process of the differential multi-head part. Because the local receptive field of traditional convolution is difficult to capture cross-regional correlation features, the differential multi-head attention mechanism can suppress irrelevant noise and model long-range dependencies between pixels through global interaction, learn more feature information to enhance the quality of reconstructed images, and support the network model detection task.

[0127] The subpixel convolution module amplifies the enhanced feature map through subpixel convolution and outputs the enhanced and amplified defect image data.

[0128] The subpixel convolution module employs an efficient upsampling method that combines defect feature extraction with channel rearrangement via convolutional layers. Its core principle is to learn the mapping from low-resolution to high-resolution features using standard convolution, and then achieve lossless upsampling of low-resolution images through information recombination from channel dimension to spatial dimension. This process consists of two steps. In the feature extraction stage, the acquired low-resolution defect enhancement feature map (C×H×W, number of channels, length, width) is processed by a standard convolutional layer (k×k). Based on the set target magnification factor r, a standard convolution with a kernel size of k×k is used to generate a feature tensor with a channel number equal to the square of the target magnification factor r from the enhanced feature map of the defect features. The resulting output feature channel number can be expanded to r. 2 C, while keeping the spatial dimension (H×W) unchanged. In the pixel reassembly stage, a periodic filtering operation maps the deep-dimensional channel information to the spatial dimension, resulting in r... 2 The C×H×W feature tensor is decomposed along the channel dimension to obtain r. 2 Each C×H×W subtensor is rearranged into a C×rH×rW high-resolution feature map as lossless upscaling of the defect image data.

[0129] ,

[0130] Among them, Input (H,W,C)The enhanced feature map represents the input defect features, where C is the number of channels and H×W is the spatial dimension. PixelShuffle() is the core operation of sub-pixel convolution, which rearranges the channel dimension information into the spatial dimension, thereby achieving image upsampling and magnification, expanding the number of channels to C×r. 2 Output (rH,rW,C) This represents the output of the defect image feature map, which is the enhanced and magnified defect reconstruction image, achieving lossless magnification and output of the low-resolution defect enhancement feature map.

[0131] The sub-pixel convolution module uses a periodic filtering operation to map deep-dimensional channel information to a spatial dimension, and then splits the resulting r-squared feature tensor according to the channel dimension to obtain r. 2 The sub-tensors are rearranged according to a chessboard pattern. The input is an H×W×C feature map, the convolution kernel is k×k, the number of channels is expanded to C×r², where r is the magnification factor, and the output is an rH×rW×C feature map. The specific arrangement is as follows:

[0132] ,

[0133] Where PS represents the PixelShuffle operation, c is the channel index, r is the upsampling factor, and T is the input tensor (H, W, C×r). 2 x and y are the spatial coordinates of the output feature map, and the modulo operation is used to locate the specific position of the sub-pixel within the block. This module can be used to perform lossless upscaling of the acquired defect image.

[0134] It also includes a loss penalty module for training the residual multi-head attention fusion model, which uses a perceptual loss function to calculate the difference in feature maps between the reconstructed defect image and the real defect image in the defect image sample, and combines adversarial loss to form a feedback mechanism to adjust the residual multi-head attention fusion model.

[0135] The perceptual loss function consists of two parts: adversarial loss and content loss. The adversarial loss, based on generative adversarial game theory, aims to enable the generator to produce realistic images, focusing more on whether the generated image closely approximates the distribution of the real image. The content loss, on the other hand, focuses on the differences in details and high-level semantic features between the generated and real images. It utilizes a feature extraction model to extract features from both the real and generated images, calculating the differences between them. This can be expressed as:

[0136] .

[0137] in The weighting factor to counteract the loss is set to 10. -3 To balance the contributions of both, the adversarial loss function minimizes the difference between the generated image and the real high-resolution image in the discriminator output, thereby enabling the generator to produce realistic images, as shown below:

[0138] .

[0139] in Let L be the real high-resolution image, G be the generator, and D be the discriminator. The generator produces a high-resolution image based on the input low-resolution image, and the discriminator needs to determine whether the generated image is generated by the generator or is real. Through continuous training, the generator can generate more realistic high-resolution images, while the discriminator can better distinguish whether an image was generated by the generator. The content loss function uses a pre-trained VGG-19 feature extraction network to extract feature maps from the reconstructed image and the real high-resolution image respectively, and calculates the difference between the two, specifically expressed as:

[0140] .

[0141] Wherein, FVGG is the feature extraction function in the VGG-19 network. The feature map is the output of the feature extraction network to show the real image of the defect. W, H, and C are the size and number of channels of the feature map, respectively. i, j, and k represent the height, width, and channel index of the feature map, respectively. i is the height index (0 to H-1), j is the width index (0 to W-1), and k is the channel index (0 to C-1).

[0142] The present invention also discloses a computer device, including a memory and a processor, wherein the memory is used to store a computer program, and the processor runs the computer program to enable the electronic device to perform the above-described method for enhancing metal surface defect images based on residual multi-head attention.

[0143] This computer device can be a server or a terminal, and its internal structure can be shown in Figure 7. The computer device includes a processor, memory, input / output interfaces (I / O), and a communication interface. The processor, memory, and I / O interfaces are connected via a system bus, and the communication interface is also connected to the system bus via the I / O interfaces. The processor provides computing and control capabilities. The memory includes non-volatile storage media and internal memory. The non-volatile storage media stores the operating system, computer programs, and a database. The internal memory provides the environment for the operation of the operating system and computer programs in the non-volatile storage media. The database stores brainwave music generation data. The I / O interfaces are used for information exchange between the processor and external devices. The communication interface is used for communication with external terminals via a network connection. When the computer program is executed by the processor, it implements a method for enhancing micro-defects on a metal surface.

[0144] Those skilled in the art will understand that the structure shown in FIG7 is merely a block diagram of a portion of the structure related to the present invention and does not constitute a limitation on the computer device to which the present invention is applied. A specific computer device may include more or fewer components than shown in the figure, or combine certain components, or have different component arrangements. In an exemplary embodiment, a computer device is provided, including a memory and a processor. The memory stores a computer program, and the processor executes the computer program to implement the steps in the above-described method embodiments.

[0145] It should be noted that the user information (including but not limited to user device information, user personal information, etc.) and data (including but not limited to data used for analysis, data stored, data displayed, etc.) involved in this invention are all information and data authorized by the user or fully authorized by all parties, and the collection, use and processing of related data must comply with relevant regulations.

[0146] Those skilled in the art will understand that all or part of the processes in the above embodiments can be implemented by a computer program instructing related hardware. The computer program can be stored in a non-volatile computer-readable storage medium. When executed, the computer program can include the processes of the embodiments described above. Any references to memory, databases, or other media used in the embodiments provided by this invention can include at least one of non-volatile and volatile memory. Non-volatile memory can include read-only memory (ROM), magnetic tape, floppy disk, flash memory, optical memory, high-density embedded non-volatile memory, resistive random access memory (ReRAM), magnetic random access memory (MRAM), ferroelectric random access memory (FRAM), phase change memory (PCM), graphene memory, etc. Volatile memory can include random access memory (RAM) or external cache memory, etc. By way of illustration and not limitation, RAM can take many forms, such as Static Random Access Memory (SRAM) or Dynamic Random Access Memory (DRAM).

[0147] The databases involved in the various embodiments provided by this invention may include at least one type of relational database and non-relational database. Non-relational databases may include, but are not limited to, blockchain-based distributed databases. The processors involved in the various embodiments provided by this invention may be general-purpose processors, central processing units, graphics processors, digital signal processors, programmable logic devices, data processing logic units, etc., and are not limited to these.

[0148] This document uses specific examples to illustrate the principles and implementation methods of the present invention. The descriptions of the above embodiments are only for the purpose of helping to understand the method and core ideas of the present invention. Furthermore, those skilled in the art will recognize that, based on the ideas of the present invention, there will be changes in the specific implementation methods and application scope. Therefore, the content of this specification should not be construed as a limitation of the present invention.

Claims

1. A method for enhancing metal surface defect images based on residual multi-head attention, characterized in that: The process includes the following steps: S1, obtaining an initial defect image containing defect features; S2, inputting the initial defect image into a pre-trained residual multi-head attention fusion model, wherein the residual multi-head attention fusion model extracts defect features from the input initial defect image, and obtains a feature sequence from the extracted defect feature map through a differential multi-head attention mechanism and converts it into an enhanced feature map of the defect features. The process of the multi-head attention mechanism is as follows: flattening the defect feature map into a one-dimensional vector, generating two sets of queries, keys, and values ​​of the defect feature map through two sets of linear transformations, segmenting the two sets of queries, keys, and values, calculating the attention score of the segmented queries, keys, and values ​​through each attention head, and concatenating the outputs of each attention head to convert it into an enhanced feature map of the defect features; S3, amplifying the enhanced feature map through sub-pixel convolution, and outputting the enhanced and amplified defect image data.

2. The method for enhancing metal surface defect images based on residual multi-head attention according to claim 1, characterized in that: The residual module of the residual multi-head attention fusion model adopts a Conv-BN-ReLU chain structure.

3. The method for enhancing metal surface defect images based on residual multi-head attention according to claim 1, characterized in that: The attention score of the attention head in step S2 is calculated using the following formula: 、 、 Where A1 and A2 are two independent projections generated from the input feature sequence, d k For each attention head dimension, Q1' and Q2' are the two sets of queries after splitting, K1' and K2' are the two sets of keys after splitting, V1' and V2' are the two sets of values ​​after splitting, T is the matrix transpose, and head i This indicates the final output of the current header. This represents a learnable differential weight parameter.

4. The method for enhancing metal surface defect images based on residual multi-head attention according to claim 1, characterized in that: Step S3 includes the following sub-steps: S31. Based on the set amplification factor r value, a standard convolution with a kernel size of k×k is used to generate a feature tensor with the number of channels equal to the square of the target amplification factor r value from the enhanced feature map of the defect features; S32. A periodic filtering operation is used to map the deep-dimensional channel information of the defect features to the spatial dimension, and the resulting feature tensor is split according to the channel dimension to obtain r. 2 Each sub-tensor is rearranged according to a chessboard pattern to obtain lossless magnified defect image data, specifically represented as follows: , where Input (H,W,C) The enhanced feature map represents the input defect features, where C is the number of channels, H×W is the spatial dimension, and the number of channels is expanded to C×r. 2 Output (rH,rW,C) This represents the feature map of the defect image output, which is the enhanced and magnified defect image data output in this step.

5. The method for enhancing metal surface defect images based on residual multi-head attention according to claim 1, characterized in that: The training process of the residual multi-head attention fusion model is as follows: K1. Obtain the defect image dataset, in which the defect image samples contain specific defect feature information and defect feature annotation text information; K2. Input the defect image samples in the defect image dataset into the residual multi-head attention fusion model, which extracts sample defect features from the input defect image samples, and obtains the sample feature sequence from the extracted sample defect feature map through the differential multi-head attention mechanism and converts it into an enhanced feature map of sample defect features; K3. Amplify the enhanced feature map of the sample defect features through sub-pixel convolution and output the enhanced and amplified defect reconstruction image; K4. Use the perceptual loss function to calculate the difference between the feature map of the defect reconstruction image and the real defect image in the defect image sample, and combine it with adversarial loss to form a feedback mechanism to adjust the residual multi-head attention fusion model; K5. Use the defect feature annotation text information in the defect image sample to evaluate and test the residual multi-head attention fusion model.

6. The method for enhancing metal surface defect images based on residual multi-head attention according to claim 5, characterized in that: In step K4, the perceptual loss function includes adversarial loss and content loss. ,in, For the perceptual loss function, For content loss function, To counteract the loss function, The weights are for the adversarial loss; the adversarial loss function utilizes the difference between the defect-reconstructed image and the real defect image in the discriminator output to prompt the generator to produce a realistic image, as shown below: Where x is the real high-resolution image, corresponding to the defect reconstruction image, L is the input low-resolution image, corresponding to the real defect image in the sample, G is the generator, G(L) is the defect image generated by the generator based on the input low-resolution image, D is the discriminator, D(x) determines the probability that x is a real image, D(G(L)) determines the probability that it is a generated image, and E[ ] represents the expectation; the content loss function uses a feature extraction network to extract feature maps of the defect reconstruction image and the real defect image respectively and calculates the difference between the two, expressed as: Where FVGG(x) is the feature map of the defect reconstruction image output by the feature extraction network. The feature map is the output of the feature extraction network to show the real image of the defect. W, H, and C are the width, height, and number of channels of the extracted feature map, respectively, and i, j, and k represent the height index, width index, and channel index within the feature map, respectively.

7. The method for enhancing metal surface defect images based on residual multi-head attention according to claim 6, characterized in that: The defect image dataset comes from the NEU-DET dataset.

8. The method for enhancing metal surface defect images based on residual multi-head attention according to claim 6, characterized in that: The feature extraction network is a pre-trained VGG-19 network model.

9. An image enhancement system based on residual multi-head attention, characterized in that: include: The image acquisition module acquires an initial defect image containing defect features; The residual multi-head attention fusion module inputs the initial defect image into a pre-trained residual multi-head attention fusion model. This model extracts defect features from the input image and uses a differential multi-head attention mechanism to obtain a feature sequence, which is then converted into an enhanced feature map of the defect features. The multi-head attention mechanism works as follows: the defect feature map is flattened into a one-dimensional vector; two sets of linear transformations are used to generate two sets of queries, keys, and values; these sets are then segmented; attention scores are calculated for each segmented query, key, and value using each attention head; and the outputs of each attention head are concatenated to convert the feature map into an enhanced feature map of the defect features. The sub-pixel convolution module amplifies the enhanced feature map using sub-pixel convolution, outputting the amplified defect image data.

10. A computer device, characterized in that, The device includes a memory and a processor, the memory being used to store a computer program, and the processor running the computer program to cause the electronic device to perform the residual multi-head attention-based image enhancement method for metal surface defects as described in any one of claims 1 to 7.

Images