An unsupervised defect segmentation algorithm based on Gaussian mixture model

Abstract

The application discloses an unsupervised defect segmentation algorithm based on a Gaussian mixture model, and comprises the following steps: S1, extracting features from a normal image to obtain a first feature map; S2, performing feature selection on the first feature map according to the variance of each channel of the first feature map, screening out topK dimensions with the maximum variance, using the features of the K dimensions to form a first new feature map, and meanwhile reserving K indexes of the K dimensions; S3, modeling the probability distribution of the features at each position on the first new feature map by using a Gaussian mixture model, and calculating model parameters; S4, extracting features from a test image in the manner of step S1 to obtain a second feature map, performing dimension screening on the second feature map according to the K indexes obtained in step S2 to obtain a second new feature map; and S5, based on the parameters calculated in step S3, calculating an abnormal score of the features at all positions on the second new feature map respectively to obtain an abnormal map of the test image.

Description

Technical Field

[0001] This invention relates to the field of computer image processing, and specifically to an unsupervised defect segmentation algorithm based on a Gaussian mixture model. Background Technology

[0002] Defect segmentation is the process of locating abnormal regions in an image. In industry, automatically locating product defects is invaluable; however, obtaining and calibrating training samples is costly, and the types of anomalies are often diverse. Traditional supervised learning methods are ineffective in this context. In recent years, unsupervised learning methods have emerged, which train models on only normal samples to enable them to identify abnormal samples. Unsupervised segmentation algorithms go a step further, enabling models not only to identify abnormal samples but also to locate abnormal regions, making these algorithms even more promising for industrial applications.

[0003] For example, Thomas Defard et al. proposed the unsupervised segmentation algorithm PaDiM based on a pre-trained model and Mahalanobis distance (Defard, Thomas; Setkov, Aleksandr; Loesch, Angelique; Audigier, Romaric (2020): PaDiM: a Patch Distribution Modeling Framework for Anomaly Detection and Localization. Available online at http: / / arxiv.org / pdf / 2011.08785v1.). The PaDiM algorithm process can be summarized as follows: Features are extracted from all normal images using a WideResNet-50 model pre-trained on ImageNet; then, the feature maps output from the Block layers 2, 3, and 4 are extracted, uniformly scaled to 56×56, and then these three feature maps are concatenated along the channel dimension. Since the number of channels in the feature map is very large at this point, 550-dimensional features are selected using a random sampling method. Assuming the feature vector at each location in the feature map originates from a multidimensional Gaussian distribution, the mean vector and covariance matrix of the Gaussian distribution at each location can be calculated using the feature maps of all normal images. When testing a new image, features are first extracted and stitched together. Then, the Mahalanobis distance at that location is calculated using the previously calculated mean and covariance; this value serves as the anomaly score at that location. Finally, this score is enlarged to the size of the original image to obtain the anomaly map for the test image. The location of defects can then be identified using this anomaly map.

[0004] While existing methods utilize intermediate features from the network and preserve multi-scale information, they lose some information during dimensionality reduction by simply randomly selecting fixed-dimensional features. Furthermore, the assumption that the vector at a certain location in the feature map follows a multidimensional Gaussian distribution is insufficient to fully describe the true distribution of the feature vectors. When testing images, only the Gaussian distribution parameters at the corresponding locations in the test image's feature map are used to calculate anomaly scores, leading to decreased accuracy and robustness when testing translated or rotated images. Summary of the Invention

[0005] The main purpose of this invention is to address the problems of excessive information loss during dimensionality reduction, insufficient modeling of distribution, and inadequate robustness in existing unsupervised defect segmentation algorithms. This invention proposes an unsupervised defect segmentation algorithm based on a Gaussian mixture model.

[0006] To achieve the above objectives, one aspect of the present invention provides the following technical solution:

[0007] An unsupervised defect segmentation algorithm based on Gaussian mixture model includes the following steps: S1, extracting features from a normal image to obtain a first feature map E. n S2, according to the first feature map E n Feature selection is performed on the variance of each channel, and the top K dimensions with the largest variance are selected. These K dimensions are then used to construct the first new feature map E. n ', while retaining the K indices of these K dimensions; S3, applying a Gaussian mixture model to the first new feature map E n Model the probability distribution of features at each location and calculate the model parameters; S4. Extract features from the test image in the manner of step S1 to obtain the second feature map E. t Then, for the second feature map E t Dimensional filtering is performed on the K indices obtained in step S2 to obtain the second new feature map E. t S5. Based on the parameters calculated in step S3, process the second new feature map E. t Anomaly scores are calculated for features at all locations to obtain an anomaly map of the test image.

[0008] In another aspect, the present invention also proposes a computer-readable storage medium having a computer program stored thereon, which, when executed by a processor, can implement the steps of the aforementioned unsupervised defect segmentation algorithm.

[0009] Compared with existing technologies, the unsupervised defect segmentation algorithm based on Gaussian Mixture Model proposed in this invention has the following technical advantages: Regarding dimension selection, a variance-based feature selection technique is used to select dimensions with larger variance, reducing information loss during dimensionality reduction; using a Gaussian Mixture Model (GMM) to model feature vectors can more accurately describe the distribution characteristics of feature vectors than a single multivariate Gaussian distribution, thus providing a more accurate anomaly score; finally, during inference, for a specific feature vector of the test image, the results calculated from multiple distributions around its corresponding location are fused, making the algorithm more robust. Attached Figure Description

[0010] Figure 1 This is a flowchart illustrating the unsupervised defect segmentation algorithm based on a Gaussian mixture model according to an embodiment of the present invention. Detailed Implementation

[0011] The present invention will be further described below with reference to the accompanying drawings and specific embodiments. This invention proposes an unsupervised defect segmentation algorithm based on a Gaussian mixture model; please refer to [reference needed]. Figure 1 The algorithm includes the following steps S1 to S5:

[0012] S1, For normal image I n Feature extraction is performed to obtain its feature map, denoted as E. n ; where I n This represents the nth normal image among N normal image samples, therefore n = 1, 2, ..., N.

[0013] In step S1, the feature generation module M1 used for feature extraction is WideResNet-50 pre-trained on ImageNet, for the nth normal image I out of N samples. n Let n = 1, 2, ..., N. After forward computation, the intermediate features of the 2nd, 3rd, and 4th blocks are obtained, denoted as f2, f3, and f4, respectively. Then, the intermediate features f2, f3, and f4 are enlarged to the scale of the largest feature map. In this example, f2 is the largest feature map, so f3 and f4 are enlarged to the same scale as f2, i.e., 56×56. The enlarged intermediate features are then concatenated along the channel dimension to obtain the feature map E. n .

[0014] S2, according to feature map E n Feature selection is performed on the variance of each channel, and the top K dimensions with the largest variance are selected. These K dimensions are then used to construct a new feature map E. n ', while retaining the K indexes of these K dimensions. For example... Figure 1As shown, module V is used to perform this step. Regarding the value of K, it should not be too small compared to the original dimension of the feature map (1792 dimensions). In experiments, it has been verified that when K=1000, not much information is lost and the performance is not significantly reduced. Therefore, it is understandable that K can be any value above 1000 within the range smaller than the original feature dimension.

[0015] In step S2, the feature map E is solved according to the following formula. n The mean μ in each channel c Sum of variance D(x) c ):

[0016]

[0017]

[0018] Where W and H are the feature maps E n width and length; x ncij x represents the value at position (i,j) of the c-th channel of the first feature map En. c Representative feature map E n The value at channel c.

[0019] This invention utilizes a variance-based feature selection method for vectors, enabling the vectors to retain more information during dimensionality reduction. Existing methods sometimes employ random sampling during dimensionality reduction, which results in the loss of more semantic information. Other dimensionality reduction algorithms are extremely time-consuming. This invention, however, takes into account that dimensions with larger variance contain more information and employs a variance-based feature selection method for dimensionality reduction, which is both fast and preserves more information, facilitating downstream tasks.

[0020] S3. Applying a Gaussian Mixture Model (GMM) to the new feature map E n The probability distribution of features at each location is modeled, and model parameters are calculated. A Gaussian mixture model is used to model the vectors, making the modeling more accurate. Existing methods assume that the feature vectors of a certain region follow a multivariate Gaussian distribution when modeling features. However, in real industrial scenarios, this assumption may not hold true. If there are large differences within normal samples, this invention uses a Gaussian mixture model to better model the data distribution.

[0021] A Gaussian Mixture Model (GMM) can be viewed as a combination of K Gaussian models (sub-models). Therefore, this step uses the GMM to analyze the new feature map E. n Modeling the probability distribution of features at each location, the resulting probability distribution can be expressed as:

[0022]

[0023] Where x represents the new feature map E n The feature vector at position (i,j) represents the i-th row and j-th column of the feature map; θ represents the parameter set of the probability distribution; P(x|θ) represents the probability density of x under the distribution represented by θ; α k It is the probability that the observed feature belongs to the k-th Gaussian model, α k ≥0, φ(x|θ k ) represents the density function of the k-th Gaussian model distribution, given by θ k Sure, μ k σ k and represent the mean and covariance of the k-th Gaussian model, respectively.

[0024] The model parameters calculated in step S3 include α k and the parameters θ of each Gaussian model k The calculation method is an iterative calculation using the Expectation-Maximum (EM) algorithm, as detailed below:

[0025] Step E: Calculate each data point x based on the current parameters. m The probability γ from the k-th Gaussian model mk :

[0026]

[0027] Where, x m φ(x) represents the m-th vector among all vectors at position (i,j). m |θ k ) represents a constant value x m The density function of the k-th Gaussian model distribution;

[0028] M-step: Calculate the model parameters for the new iteration.

[0029]

[0030]

[0031]

[0032] Repeat steps E and M above until convergence (||θ) s+1 -θ s ||<ε); ε is a very small constant, for example, in or around the range of 0.001 to 0.0001. Where θ s This represents the parameters of all K Gaussian models at the s-th iteration.

[0033] For the new feature map En The vectors at all locations are calculated using the E-step and M-step methods described above to obtain the new feature map E. n The set of Gaussian distribution parameters at all locations {Θ ij |1≤i≤W,1≤j≤H}.

[0034] S4. Test image I for abnormalities t Feature extraction is performed according to step S1 to obtain the corresponding feature map, denoted as E. t Then, for feature map E t Dimensional filtering is performed using the K indices obtained in step S2 to obtain a new feature map E of the test image. t '.

[0035] S5. Based on the parameters finally calculated in step S3, a new feature map E of the test image is generated. t Anomaly scores are calculated for features at all locations to obtain an anomaly map of the test image.

[0036] In step S5, for each position (i,j) of the new feature map Et', the feature vector x ij The parameters of the Gaussian mixture model at the corresponding location are Θ. ij And the parameter Θ of its surrounding location. i-1，j-1 Θ i-1，j Θ i-1，j+1 Θ i，j-1 Θ i，j+1 Θ i+1，j-1 Θ i+1，j Θ i+1，j+1 For Θ ij Θ i-1，j-1 Θ i-1，j Θ i-1，j+1 Θ i，j-1 Θ i，j+1 Θ i+1，j-1 Θ i+1，j Θ i+1，j+1 These parameters are calculated using the following formula for vector x. ij The probability density under the distributions represented by these parameters:

[0037]

[0038] Wherein, P(x ij |Θ) represents the vector x ij The probability density function φ(x) under the distribution represented by parameter Θ. ij |Θ k ) represents vector x ij The probability density function under the distribution of each Gaussian component. It should be understood that the above formula is used to calculate the vector x. ijIn Θ ij Θ i-1，j-1 Θ i-1，j Θ i-1，j+1 Θ i，j-1 Θ i，j+1 Θ i+1，j-1 Θ i+1，j Θ i+1，j+1 When these parameters represent the probability density of the distribution, simply replace Θ in the formula with the corresponding specific parameters.

[0039] Ultimately, vector x ij The abnormal score is

[0040] score(x ij )=min(P(x ij |Θ))

[0041] Solve for vector x ij The method of anomaly scoring can be called domain information capture. Each vector x ij Not only using the model Θ at the same location ij When calculating outlier scores, Θ is also considered. ij The method considers neighboring vectors; finally, the minimum of all calculated results is taken as the final anomaly score, making the anomaly score estimation more accurate. Existing methods, when detecting a given feature vector, only consider the Gaussian model at the corresponding location. This approach is difficult to handle when the test image moves or is deflected, resulting in poor robustness. This invention integrates the results of multiple models around the corresponding location of the feature vector, enabling the model to capture more local information and increasing robustness.

[0042] The feature vectors x at all positions (i,j) are obtained. ij After determining the anomaly score, the anomaly map of the test image is obtained.

[0043] Another embodiment of the present invention provides a computer-readable storage medium storing a computer program thereon, which, when executed by a processor, implements the steps of the unsupervised defect segmentation algorithm described in the foregoing embodiments. Based on this understanding, the technical solution of the present invention can be embodied in the form of a software product, which can be stored in a non-volatile storage medium (such as a CD-ROM, USB flash drive, or portable hard drive), including several instructions to cause a computer device (such as a personal computer, server, or network device, etc.) to execute the method steps of the various embodiments of this application.

[0044] The above description, in conjunction with specific preferred embodiments, provides a further detailed explanation of the present invention. It should not be construed that the specific implementation of the present invention is limited to these descriptions. For those skilled in the art, several equivalent substitutions or obvious modifications can be made without departing from the concept of the present invention, and all such modifications, achieving the same performance or purpose, should be considered within the scope of protection of the present invention.

Claims

1. An unsupervised defect segmentation algorithm based on a Gaussian mixture model, characterized in that, Includes the following steps: S1. Extract features from the normal image to obtain the first feature map E. n ; S2, according to the first feature map E n Feature selection is performed on the variance of each channel, and the top K dimensions with the largest variance are selected. These K dimensions are then used to construct the first new feature map E. n ', while retaining K indexes for these K dimensions; S3. Apply a Gaussian mixture model to the first new feature map E. n Model the probability distribution of features at each location and calculate the model parameters; S4. Perform feature extraction on the test image in the same manner as in step S1 to obtain the second feature map E. t Then, for the second feature map E t Dimensional filtering is performed on the K indices obtained in step S2 to obtain the second new feature map E. t '; S5. Based on the parameters calculated in step S3, the second new feature map E... t Anomaly scores are calculated for features at all locations to obtain an anomaly map of the test image.

2. The unsupervised defect segmentation algorithm based on Gaussian mixture model as described in claim 1, characterized in that, Step S1 specifically includes: S11. Using Wide ResNet-50 pre-trained on ImageNet as the feature extractor, for the nth normal image I out of N samples. n For n = 1, 2, ..., N, after forward computation, the intermediate features of the 2nd, 3rd, and 4th blocks are obtained, which are denoted as f2, f3, and f4 respectively; S12. Magnify the intermediate features f2, f3, and f4 to a preset scale, and then stitch the magnified intermediate features together along the channel dimension to obtain the first feature map E. n .

3. The unsupervised defect segmentation algorithm based on Gaussian mixture model as described in claim 2, characterized in that, Step S12, which involves enlarging the intermediate features f2, f3, and f4 to a preset scale, includes: using the largest feature map scale among the three intermediate features f2, f3, and f4 as the preset scale, and enlarging the other two to the preset scale.

4. The unsupervised defect segmentation algorithm based on Gaussian mixture model as described in claim 1, characterized in that, In step S2, the first feature map E is solved in the following way. n The mean μ in each channel c Sum of variance D(x) c ): Where N is the total number of normal image samples, and W and H are the first feature map E n width and length; x ncij x represents the value at position (i,j) of the c-th channel of the first feature map En. c Representing the first feature map E n The value at channel c.

5. The unsupervised defect segmentation algorithm based on Gaussian mixture model as described in claim 1, characterized in that, The Gaussian mixture model described in step S3 is composed of K Gaussian models, and the corresponding probability distribution is as follows: Where x represents the first new feature map E n The eigenvector at position (i,j) is given by θ, which represents the set of parameters for the probability distribution. P(x|θ) represents the probability density of x under the distribution represented by θ. k It is the probability that the observed feature belongs to the k-th Gaussian model, α k ≥0, φ(x|θ k ) represents the density function of the k-th Gaussian model distribution, given by θ k Sure, μ k σ k and represent the mean and covariance of the k-th Gaussian model, respectively; The model parameters calculated in step S3 include α k and the parameters θ of each Gaussian model k The calculation method is to perform iterative calculation using the expectation-maximization algorithm.

6. The unsupervised defect segmentation algorithm based on Gaussian mixture model as described in claim 5, characterized in that, Step S3 iteratively calculates α using the expectation-maximization algorithm. k and θ k The steps include: Step E: Calculate each data point x based on the current parameters. m The probability γ from the k-th Gaussian model mk : Where, x m φ(x) represents the m-th vector among all vectors at position (i,j), where N is the total number of normal image samples. m |θ k ) represents a constant value x m The density function of the k-th Gaussian model distribution; M-step: Calculate the model parameters for the new iteration. Repeat steps E and M above until convergence (||θ) s+1 -θ s ||<ε);where θ s This represents the parameters of all K Gaussian models at the s-th iteration; For the first new feature map E n Perform the above E-step and M-step calculations on the vectors at all locations to obtain the set of Gaussian distribution parameters {Θ} at all locations. ij |1≤i≤W,1≤j≤H}.

7. The unsupervised defect segmentation algorithm based on Gaussian mixture model as described in claim 6, characterized in that, Step S5 specifically includes: For the second new feature map E t 'Feature vector x at each position (i,j) ij The parameters of the Gaussian mixture model at the corresponding location are Θ. ij And the parameter Θ of its surrounding location. i-1，j-1 Θ i-1，j Θ i-1，j+1 Θ i，j-1 Θ i，j+1 Θ i+1，j-1 Θ i+1，j Θ i+1，j+1 ; For each parameter Θ ij Θ i-1，j-1 Θ i-1，j Θ i-1，j+1 Θ i，j-1 Θ i，j+1 Θ i+1，j-1 Θ i+1，j Θ i+1，j+1 The vector x is calculated using the following formula. ij The probability density under the distributions represented by these parameters: Wherein, P(x ij |Θ) represents the vector x ij The probability density function φ(x) under the distribution represented by parameter Θ. ij |Θ k ) represents vector x ij The probability density under the distribution of each Gaussian component; Ultimately, vector x ij The abnormal score is score(x ij )=min(P(x ij |Θ)) Obtain the feature vector x at all positions (i,j) ij After determining the abnormal score, the abnormal map of the test image is obtained.

8. A computer-readable storage medium having a computer program stored thereon, characterized in that, When executed by a processor, the computer program implements the steps of the unsupervised defect segmentation algorithm as described in any one of claims 1-7.

Images