An industrial image repairing and feature learning method based on an attention mechanism

Abstract

The application discloses an industrial image repairing and feature learning method based on an attention mechanism, and particularly relates to the technical field of computer vision and industrial image repairing. The application comprises an overall network structure, and the overall network structure is based on a Swin-Transformer framework. The overall network structure comprises a shallow feature extraction module, a UNet feature extraction module, a reconstruction module, a Loss calculation module and a feature learning module. The industrial image repairing and feature learning method based on the attention mechanism uses a Swin-Transformer Block as a basic module, and then the overall network structure is stacked in the form of a UNet architecture, namely a Swin-Transformer UNet (STUnet), which is used for learning efficient and robust high-dimensional image features, so that image information repairing is effectively realized. Thanks to the network design, the STUnet and the training process can learn high-quality robust features, and have strong generality. The features can be directly used for subsequent downstream quality inspection tasks, such as common industrial quality inspection visual tasks, for example, classification, measurement, positioning and the like.

Description

TECHNICAL FIELD

[0001] The present application relates to the technical field of computer vision and industrial image restoration, in particular to an industrial image restoration and feature learning method based on an attention mechanism. BACKGROUND

[0002] During the acquisition process of industrial images, unnecessary information blurring or information loss may be caused by environmental interference, sensor sampling anomalies and other factors. These anomalies may affect image quality, making it difficult for humans or computers to accurately identify the image. In industrial production and quality inspection links, images are often used for diagnosis, detection, monitoring and other applications. Therefore, sampling quality problems of industrial images may lead to misjudgment, false diagnosis and other serious consequences.

[0003] Commonly used sampling images for industrial quality inspection tasks may have the following quality problems due to factors such as sampling environment and equipment aging: noise, blurring, partial missing of the picture, color cast and uneven lighting. For the above common image quality problems, obvious interference may be caused in the subsequent quality inspection link, and even the production quality and production safety may be seriously affected. Although some computer-based image denoising and smoothing models exist at the present stage, which can repair image information to a certain extent, these models also have obvious shortcomings, making them unable to be widely used in actual production and quality inspection. Therefore, industrial image restoration is very important, which can improve the quality of images, thereby further improving the accuracy, efficiency and safety of quality inspection.

[0004] Among them, the traditional image restoration method is mainly based on the image smoothing method in traditional computer vision, mainly using a filter to filter the specific frequency domain and space in the image. Although this method is simple in algorithm, the effect is limited, especially when facing complex image information quality anomalies, the effect is often unsatisfactory. At the same time, the global filtering function usually causes the loss of image details, thereby seriously affecting the accuracy of subsequent quality inspection tasks. Therefore, in recent years, some deep learning models can learn the characteristics of noise in a specific scene by training a large number of noisy images and corresponding denoising images, thereby achieving effective image information restoration. However, although the above method has certain performance advantages compared with traditional machine vision models, it also has serious generalization problems, that is, the model is difficult to adapt to different complex industrial scenes. SUMMARY

[0005] (I) Technical problems solved

[0006] In view of the deficiencies of the prior art, the present application provides an industrial image restoration and feature learning method based on an attention mechanism, which solves the problem that existing industrial image restoration is limited by model accuracy, generalization and robustness, making it difficult to be simply reused in the field of industrial image denoising and industrial quality inspection.

[0007] (ii) Technical Solution

[0008] To achieve the above object, the present application is implemented by the following technical solution: an industrial image restoration based on an attention mechanism, comprising an overall network structure, and the overall network structure is based on a Swin-Transformer framework, the overall network structure comprises a shallow feature extraction module, a UNet feature extraction module, a reconstruction module, a Loss calculation module and a feature learning module, and the industrial image restoration specifically comprises the following steps:

[0009] Step one, the initial image is first passed through the shallow feature extraction module to obtain shallow features;

[0010] Step two, the extracted shallow features are input to the UNet feature extraction module to extract high-dimensional features and perform information reconstruction;

[0011] Step three, the final generated image is obtained by using the reconstruction module.

[0012] Preferably, the shallow feature extraction module inputs an image I∈R H×W×3 Here, the defect refers to the case that the image noise is obvious, the image is blurred, the image is partially lost, and the image is overexposed, H and W represent the resolution of the image respectively, and a 3×3 convolution layer is used to obtain shallow feature information F shallow ∈R H×W×C , wherein C represents the number of convolution kernel groups, and the specific formula is as follows: F shallow =M SFE (I).

[0013] Preferably, the shallow feature information F shallow is extracted by the UNet feature extraction module to obtain deep features F st ∈R H×W×C , and the specific formula is as follows: F st =M UFE (F shallow ).

[0014] Preferably, the M UFE () module adopts a UNet structure, and a total of 7 Swin-Transformer Block (STB) units are stacked in series to replace the convolution module in the traditional UNet framework, and at the same time, in order to ensure the stability of model training and feature learning, the STB in the UFE adopts a skip connection mode for feature transmission, and after obtaining the deep features F st , a 3×3 convolution layer M RM () is connected at the end of the network to obtain the final up-sampling output The specific formula is as follows:

[0015] Preferably, in order to obtain effective depth features F st , M UFE Each Swin-TransformerBlock in the () module contains 8 Swin-Transformer Layers (STLs), and each STL is connected in series with 2 different attention mechanisms, namely Window Multi-head Self-Attention (WMSA) and Shifted-Window Multi-head Self-Attention (SWMSA). The specific working mode of the network is as follows:

[0016] f L WMSA = WMSA (LN(f L-1 ))+f L-1

[0017] f L MLP = MLP (LN(f L WMSA ))+f L WMSA

[0018] f L SWMSA = SWMSA (LN(f L MLP ))+f L MLP

[0019] f L = MLP (LN(f L SWMSA ))+f L SWMSA

[0020] wherein the input f L-1 of each STB, after the Layer Norm module, WMSA, MLP, SWMSA and other modules, finally obtains the output f L of the STB. 8 STL modules are stacked in series in each STB. For each MSA module, its internal mathematical structure is equivalent to the standard multi-head attention mechanism, i.e.:

[0021] Q, K, V = Conv (x)

[0022]

[0023] y = Conv (Attention * V)

[0024] Wherein, x, y respectively represent the input and output of the current MSA module, Conv represents the convolution operation, Q, K, V respectively represent three attention matrices calculated, d k Indicates the dimension.

[0025] Preferably, Loss considers information restoration degree SE, generated image signal-to-noise ratio PSNR, and generated image structural similarity SSIM, therefore, the formula definition of the calculation module of Loss is as follows:

[0026]

[0027] Wherein, α, β, γ are penalty factors, and the specific definitions of SSE, PSNR and SSIM are as follows:

[0028]

[0029] Wherein, Y i And Respectively represent the i-th original image and the corresponding image inpainting result, N represents batchsize, and μ Y And Respectively represent the mean of the real image and the predicted image, And Respectively represent the variance of the real image and the predicted image, Is the covariance of the real image and the predicted image, c1 and c2 are constants, and are usually set to (0.01*255) 2 And (0.03*255) 2 .

[0030] The application also provides an industrial image inpainting feature learning method based on an attention mechanism, which can introduce other downstream tasks while training, and simultaneously train to ensure the generality and robustness of the extracted features, then Loss can be generalized as:

[0031]

[0032] Wherein, L downstream Can be in any form, for example, in a classification task, it can be Cross Entropy Loss, in a positioning or measurement task, it can be MSE, etc., and the specific steps include:

[0033] S1, select a Swin-Transformer framework pre-trained on ImageNet-22K as a benchmark model;

[0034] S2, make a dataset, that is, use ImageNet-1K dataset, randomly add white noise, overexposed points, and pixel random loss to generate adversarial samples;

[0035] S3, in the STUNet model, the feature learning module is supplemented as a fully connected network layer and a Softmax layer, for image classification;

[0036] S4, on the basis of the model of S1, the initial STUNet is fine-tuned by using the data set obtained in S2, and meanwhile, the feature learning module in S3 is trained by using the image class labels of ImageNet-1K, so that the features extracted by the STUNet have generality, and at this time, the overall loss of the STUNet network is defined as:

[0037]

[0038] Wherein, L CrossEntropy =-(t i log(p i )+(1-t i )log(1-p i )),t i ,p i respectively correspond to the true value and the predicted probability of the label;

[0039] S5, an industrial image data set is made for the secondary fine-tuning of the STUNet model, so as to adapt to the inconsistency of the industrial image and the data style scene of ImageNet-1K, and the data set should be made by adopting the same standard and steps as S2;

[0040] S6, the model in S4 is fine-tuned by using the data set obtained in S5, and here, considering the diversity of the industrial quality inspection task, the feature learning module can be replaced in a targeted manner, for example, in the classification quality inspection task, it can be a fully connected layer and a Softmax layer, and in the positioning or measurement task, it can be a fully connected layer.

[0041] Advantages

[0042] The application provides an industrial image repair and feature learning method based on an attention mechanism.

[0043] 1、The industrial image repair and feature learning method based on the attention mechanism, by using Swin-Transformer Block as a basic module, then the whole is stacked in the form of UNet architecture, namely Swin-Transformer UNet (STUnet), for learning efficient and robust high-dimensional image features, so as to effectively realize image information repair.

[0044] 2、The industrial image repair and feature learning method based on the attention mechanism can learn high-quality robust features by benefiting from the network design and the training process, has strong universality, and the features can be directly used for subsequent downstream quality inspection tasks, such as classification, measurement, positioning and other common industrial quality inspection visual tasks. BRIEF DESCRIPTION OF DRAWINGS

[0045] Figure 1 The figure is a schematic diagram of the overall network framework of the application;

[0046] Figure 2 The figure is a schematic diagram of the STB and STL of the application;

[0047] Figure 3 The figure is an example 1 of the image repair result of the algorithm of the application;

[0048] Figure 4 The figure is an example 2 of the image repair result of the algorithm of the application. DETAILED DESCRIPTION

[0049] The technical solutions in the embodiments of the application are described clearly and completely, obviously, the described embodiments are only part of the embodiments of the application, not all the embodiments. Based on the embodiments in the application, all other embodiments obtained by those skilled in the art without creative labor are within the scope of protection of the application.

[0050] Please refer to Figures 1-4 The embodiment of the application provides a technical solution: an industrial image repair based on an attention mechanism, including an overall network structure, and the overall network structure is based on a Swin-Transformer framework, the overall network structure includes a shallow feature extraction module, a UNet feature extraction module, a reconstruction module, a Loss calculation module and a feature learning module, the industrial image repair specifically includes the following steps: step one, the initial image is first passed through the shallow feature extraction module to obtain shallow features; step two, the extracted shallow features are input to the UNet feature extraction module to extract high-dimensional features and perform information reconstruction; step three, the final generated image is obtained by using the reconstruction module. Use Swin-Transformer Block as the basic module, then stack the whole in the form of UNet architecture, that is, Swin-Transformer UNet (STUnet), which is used to learn efficient and robust high-dimensional image features, thereby effectively realizing image information repair.

[0051] Please refer to Figure 1 and 2 The shallow feature extraction module inputs the flawed image I∈R H×W×3Here, the flaw refers to the image noise, image blur, image local loss, and image overexposure, H and W represent the resolution of the image, and a 3x3 convolution layer is used to obtain the shallow feature information F shallow ∈R H×W×C where C represents the number of convolution kernel groups, and the specific formula is as follows: F shallow = M SFE (I).

[0052] Preferably, the shallow feature information F shallow The deep feature F st ∈R H×W×C , and the specific formula is as follows: F st = M UFE (F shallow ).

[0053] Preferably, the M UFE () module adopts the UNet structure, and a total of 7 Swin-Transformer Block (STB) units are stacked in series to replace the convolution module in the traditional UNet framework. In order to ensure the stability of model training and feature learning, the STB in the UFE adopts a skip connection mode for feature transmission. After obtaining the deep feature F st , the network is finally connected to a 3x3 convolution layer M RM (), which is used to obtain the final up-sampling output The specific formula is as follows:

[0054] Preferably, in order to obtain effective deep feature F st , each Swin-Transformer Block in the M UFE () module contains 8 Swin-Transformer Layers (STL), and each STL adopts a series connection mode to connect 2 different attention mechanisms, namely Window Multi-head Self-Attention (WMSA) and Shifted-Window Multi-head Self-Attention (SWMSA). The specific working mode of the network is as follows:

[0055] f L WMSA = WMSA(LN(f L-1 ))+f L-1

[0056] f L MLP = MLP(LN(f LWMSA ))+f L WMSA

[0057] f L SWMSA = SWMSA(LN(f L MLP ))+f L MLP

[0058] f L = MLP(LN(f L SWMSA ))+f L SWMSA

[0059] wherein the input f L-1 of each STB, after passing through the Layer Norm module, WMSA, MLP, SWMSA and other modules, finally obtains the output f L of the STB, 8 STL modules are stacked in series in each STB, and for each MSA module, the internal mathematical structure is equivalent to the standard multi-head attention mechanism, that is:

[0060] Q,K,V = Conv(x)

[0061]

[0062] y = Conv(Attention * V)

[0063] wherein x, y respectively represent the input and output of the current MSA module, Conv represents the convolution operation, Q, K, V respectively represent the three attention matrices calculated, d k represents the dimension.

[0064] Preferably, the Loss considers the information restoration degree SE, the generated image signal-to-noise ratio PSNR, and the generated image structural similarity SSIM, therefore, the calculation module formula of the Loss is defined as follows:

[0065]

[0066] wherein α, β, γ are penalty factors, and the specific definitions of SSE, PSNR and SSIM are as follows:

[0067]

[0068] wherein Y i and respectively represent the ith original image and the corresponding image inpainting result, N represents the batch size, and μ Y and are the mean of the real and predicted images respectively, and are the variance of the real and predicted images respectively, is the covariance of the real and predicted images, c1 and c2 are constants, usually set to (0.01 x 255) 2 and (0.03 x 255) 2 .

[0069] Referring to Figures 1-4 , the application also provides an industrial image repair feature learning method based on an attention mechanism, which can introduce other downstream tasks while training, and synchronously train to ensure the universality and robustness of the extracted features, so that Loss can be generalized as:

[0070]

[0071] where L downstream can be of any form, for example, in a classification task, it can be Cross Entropy Loss, in a positioning or measurement task, it can be MSE, etc., and specifically includes the following steps: S1, selecting a Swin-Transformer framework pre-trained on ImageNet-22K as a benchmark model; S2, making a dataset, i.e., using an ImageNet-1K dataset, randomly adding white noise, overexposed points, and pixel random loss to generate adversarial samples; S3, in the STUNet model, supplementing the feature learning module with a fully connected network layer and a Softmax layer for image classification; S4, based on the model of S1, fine-tune the initial STUNet using the dataset obtained in S2, and at the same time, use the image class labels of ImageNet-1K to train the feature learning module in S3, so as to ensure that the features extracted by STUNet have general applicability, at this time the overall Loss of STUNet network is defined as:

[0072]

[0073] where L CrossEntropy = -(t i log(p i )+(1-t i )log(1-p i )),t i , p irespectively corresponding to the true value and the predicted probability of the label; S5, making an industrial image dataset for the second fine-tuning of the STUNet model to adapt to the inconsistency of the industrial image and the ImageNet-1K data style scene, and the making of the dataset should adopt the same standard and steps as S2;

[0074] S6, fine-tuning the model in S4 with the dataset obtained in S5, and considering the diversity of the industrial quality inspection task, the feature learning module can be replaced accordingly, for example, in the classification quality inspection task, it can be a fully connected layer and a Softmax layer, and in the positioning or measurement task, it can be a fully connected layer. Benefiting from the network design, the STUNet and the training process can learn high-quality robust features, which have strong generality; the features can be directly used for subsequent downstream quality inspection tasks, such as classification, measurement, positioning and other common industrial quality inspection visual tasks.

[0075] Meanwhile, the contents not described in detail in the specification all belong to the prior art known to those skilled in the art.

Claims

1. An attention-based feature learning method for industrial image inpainting, comprising an overall network structure, wherein the overall network structure is based on the Swin-Transformer framework, characterized in that: The overall network structure includes a shallow feature extraction module, a UNet feature extraction module, a reconstruction module, a loss calculation module, and a feature learning module. Industrial image restoration specifically includes the following steps: Step 1: The initial image is first processed through the shallow feature extraction module to obtain shallow features; Step 2: The extracted shallow features are input into the UNet feature extraction module to extract high-dimensional features and reconstruct information; Step 3: Use the reconstruction module to obtain the final generated image; The Loss calculation takes into account the information restoration degree (SE), the signal-to-noise ratio (PSNR) of the generated image, and the structural similarity (SSIM) of the generated image. The formula for calculating the Loss is defined as follows: Where α, β, and γ are penalty factors, and the specific definitions of SE, PSNR, and SSIM are as follows: Among them, Y i as well as These represent the i-th original image and the corresponding image restoration result, respectively, where N represents the batch size and μ... Y and These are the means of the real image and the predicted image, respectively. and These are the variances of the real image and the predicted image, respectively. Let c1 and c2 be the covariance between the real image and the predicted image, and let c1 and c2 be constants, denoted as (0.01 × 255). 2 and (0.03×255) 2 .

2. The industrial image restoration feature learning method based on an attention mechanism according to claim 1, characterized in that: The shallow feature extraction module inputs a noisy image I∈R H×W×3 H and W represent the image resolution, respectively. A 3×3 convolutional layer is used to obtain image features, ultimately yielding shallow feature information F. shallow ∈R H×W×C Where C represents the number of convolutional kernel groups, the specific formula is as follows: F shallow =M SFE (I).

3. The industrial image restoration feature learning method based on an attention mechanism according to claim 2, characterized in that: The shallow feature information F shallow Deep features F are extracted using the UNet feature extraction module. st ∈R H×W×C The specific formula is as follows: F st =M UFE (F shallow ).

4. The industrial image restoration feature learning method based on an attention mechanism according to claim 3, characterized in that: The UNet feature extraction module adopts the UNet structure, which mainly uses the Swin-Transformer Block. Each Swin-Transformer Block contains 8 Swin-Transformer Layers, and each Swin-Transformer Layer contains two different attention mechanisms: window multi-head self-attention and shifted-window multi-head self-attention. This is used to obtain the deep features F. st Finally, the network is connected to a 3×3 convolutional layer M. RM (), used to obtain the final upsampled output. The specific working mode of the network is as follows: f L WMSA =WMSA(LN(f L-1 ))+f L-1 f L MLP =MLP(LN(f L WMSA ))+f L WMSA f L SWMSA =SWMSA(LN(f L MLP ))+f L MLP f L =MLP(LN(f L SWMSA ))+f L SWMSA 5. The industrial image restoration feature learning method based on an attention mechanism according to claim 4, characterized in that: The input f of each STB L-1 After passing through the Layer Norm module, WMSA, MLP, and SWMSA modules, the final output f of the STB is obtained. L For each MSA module, its internal mathematical construction is equivalent to the standard multi-head attention mechanism, that is: Q,K,V=Conv(x) y = Conv(Attention * V) Where x and y represent the input and output of the current MSA module, respectively, Conv represents the convolution operation, Q, K, and V represent the three attention matrices computed, and d k Indicates dimension.

6. According to any one of claims 1-5, in the industrial image restoration feature learning method based on an attention mechanism, other downstream tasks are introduced during training, and training is performed synchronously to ensure the universality and robustness of the extracted features. The Loss generalization is then: Where L downstream It can take any form; in classification tasks it is Cross Entropy Loss, and in localization or measurement tasks it is MSE. Specifically, it includes the following steps: S1. Select the Swin-Transformer framework pre-trained on ImageNet-22K as the baseline model; S2. Create a dataset, i.e., use the ImageNet-1K dataset to generate adversarial samples by randomly adding white noise, overexposure points, and random pixel loss to each of the images. S3. In the STUNet model, the feature learning module is supplemented with a fully connected network layer and a Softmax layer for image classification. S4. Based on the model described in S1, the initial STUNet is fine-tuned using the dataset obtained in S2. Simultaneously, the feature learning module in S3 is trained using ImageNet-1K image category labels to ensure the generalizability of the features extracted by STUNet. At this point, the overall loss of the STUNet network is defined as: Among them, L CrossEntropy =-(t) i log(p i )+(1-t i log(1-p) i )), t i p i These correspond to the actual value and predicted probability of the label, respectively. S5. Create an industrial image dataset for secondary fine-tuning of the STUNet model to address the inconsistency between industrial images and ImageNet-1K data style scenarios. The dataset is created using the same standards and steps as in S2. S6. Fine-tune the model in S4 using the dataset obtained in S5. Considering the diversity of industrial quality inspection tasks, the feature learning module is replaced accordingly. In classification quality inspection tasks, it is a fully connected layer and a Softmax layer, while in localization or measurement tasks, it is a fully connected layer.

Images