An image anomaly detection method based on a text-to-image diffusion model

Abstract

This invention discloses an image anomaly detection method based on a text-to-image diffusion model. It utilizes a pre-trained text-to-image diffusion model to extract semantic representations of images, and simultaneously constructs an implicit description generator to generate text embeddings describing the images, which are then injected into the text-to-image diffusion model to enhance the semantic representations. A compression module is then used to compress the semantic representations to obtain the final image features. In the detection phase, the distance between the features of the image to be detected and the normal image features used during training is calculated as the basis for anomaly detection. Experimental results show that this invention achieves good results in image anomaly detection.

Description

Technical Field

[0001] This invention relates to the field of machine learning, and specifically to an image anomaly detection method based on a text-to-image diffusion model. Background Technology

[0002] Image anomaly detection is an important research area in machine learning. It uses a dataset of normal images to build a detection model to identify anomalous images that do not match normal images in a given dataset. Anomalous images usually differ significantly from most normal images and are difficult to predict in advance. Image anomaly detection has wide applications in real life. For example, in the medical field, in tumor detection tasks, MRI images of the brains of normal individuals are readily available, but the number of tumor patients is small and the tumor conditions are complex. Therefore, image anomaly detection analyzes and models existing MRI images of the brains of normal individuals to detect the anomaly images of the brain to be tested.

[0003] In recent years, deep learning neural networks have been widely used in image anomaly detection tasks and have achieved excellent results. Existing methods can be broadly categorized into reconstruction error-based methods and distance metric-based methods. Reconstruction error-based methods encode and decode normal images, training the neural network with the reconstruction input as the target to learn the distribution patterns of normal images. In the detection phase, anomaly detection is performed by analyzing the differences between the reconstructed and reconstructed images. This type of method is sensitive to noise; slight perturbations or noise can lead to high reconstruction errors. Existing methods primarily employ distance metric-based methods. These methods extract features from the image, making the feature distribution of normal images as compact as possible. In the testing phase, the distance between the features of the test image and normal images is calculated to determine whether the image is an anomaly. This type of method requires the model to have strong representation learning capabilities to reasonably represent rare and unpredictable anomaly features in the anomaly image. The model's representation learning capability heavily depends on the type and quantity of training samples, which greatly affects the model's performance in image anomaly detection. Therefore, distance metric-based methods require further research. Summary of the Invention

[0004] The purpose of this invention is to provide an image anomaly detection method based on a text-to-image diffusion model. The method utilizes a pre-trained text-to-image diffusion model to obtain the semantic representation of the image. At the same time, an implicit description generator is proposed to generate the corresponding text embedding of the input image and inject it into the diffusion model to enhance the semantic representation of the image. Finally, a compression module is used to compress the semantic representation to obtain the final features of the image, and these features are used for distance measurement to achieve image anomaly detection.

[0005] The technical solution for implementing this invention is: an image anomaly detection method based on a pre-trained diffusion model, comprising the following steps:

[0006] Step 1: From the publicly available ImageNet dataset, select 30 classes for image anomaly detection; select one class as the normal class and the remaining 29 classes as the anomaly class; during network training, only the normal class images are used as the training dataset, and during testing, all 30 classes of images are used as the test dataset; crop all images to a uniform size of 256×256, and perform strong data augmentation and weak data augmentation on the cropped training dataset to obtain augmented images. Then construct an enhanced image training dataset {D}; crop the test dataset to a size of 256×256 without performing any other data augmentation operations, and proceed to step 2.

[0007] Step 2: Construct a global image anomaly detection network:

[0008] The overall image anomaly detection network includes a text-to-image diffusion model, an implicit description generator, and a compression module; it enhances the image... The implicit description generator is used to obtain the text embedding e, which is then used as a condition to enhance the image. The common input image diffusion model is used to obtain the semantic representation r of the image. The semantic representation is compressed to obtain image features, and then the process proceeds to step 3.

[0009] Step 3: Construct the loss function, which includes feature contrast loss and strong enhancement classification loss. The specific steps are as follows:

[0010] Step 3.1: Construct the feature contrast loss:

[0011] A contrastive learning method is introduced to train the network. In contrastive learning, similar images are called positive pairs, and dissimilar images are called negative pairs. The feature contrastive loss function is L. sim The expression is as follows:

[0012]

[0013] In the feature contrast loss function L sim In the training dataset {D}, for any augmented image... {X +}express Positive examples for the dataset, It is a set {X + Any member in}; multiple enhanced images obtained by performing multiple weak enhancements on the same strongly enhanced image are considered as positive pairs; {X -}express Negative examples of the dataset, It is a set {X - Any member in}, the sim function is used to calculate the similarity between two features, z(·) represents the output of the feature contrast layer, and τ is the temperature of the feature contrast loss, which determines the attention weight of the contrast loss function on positive and negative pairs.

[0014] Step 3.3: Construct the strongly enhanced classification loss L cls :

[0015] For the strong enhancement classification layer, the cross-entropy loss function is used to calculate the loss. Its input is the output CLS(·) of the strong enhancement classification layer and the i-th strong enhancement type actually used in the enhanced image. The specific expression is as follows:

[0016]

[0017] Step 3.4, the total loss function L is as follows:

[0018] L = L sim +α×L cls (10)

[0019] Where α is L cls The weight.

[0020] Proceed to step 4.

[0021] Step 4: Train the entire anomaly detection network using the augmented image training dataset:

[0022] The anomaly detection network is trained using the enhanced image training dataset. Throughout the process, the parameters of the CLIP model and the UNet network are fixed. Only the parameters of the linear layer and the compression module of the implicit description generator need to be trained and updated to obtain the trained anomaly detection network. Proceed to step 5.

[0023] Step 5: Input the test dataset into the trained anomaly detection network, calculate the similarity of features between the test dataset and all normal class images in the augmented image training dataset, and select the value with the highest similarity as the anomaly score of the image. For the test image y in the test dataset, its score(y) is as follows:

[0024]

[0025] Anomaly scores are used to represent the degree of anomaly in test images. The lower the score, the more likely the test image is to be abnormal, and the higher the feature similarity between similar images in the network model.

[0026] Compared with the prior art, the advantages of this invention are:

[0027] (1) This invention proposes an image anomaly detection method based on a text-to-image diffusion model, which uses a pre-trained text-to-image diffusion model to extract the semantic representation of the image, and has better generalization ability and robustness.

[0028] (2) Since there is no textual description information for images in image anomaly detection tasks, this invention designs an implicit description generator to generate image-related text embeddings based on the image. It replaces the textual description in the text-to-image diffusion model to enhance the semantic representation extracted from the text-to-image diffusion model. Attached Figure Description

[0029] Figure 1 This is a schematic diagram of the overall structure of the present invention.

[0030] Figure 2 This is a schematic diagram of the text-to-image diffusion model of the present invention.

[0031] Figure 3 This is a schematic diagram of the implicit description generator of the present invention.

[0032] Figure 4 This is a schematic diagram of the compression module of the present invention.

[0033] Figure 5 This is a schematic diagram of positive and negative example pairs of the present invention. Detailed Implementation

[0034] To make the objectives, technical solutions, and advantages of the present invention clearer, the embodiments of the present invention will be described in further detail below.

[0035] Combination Figures 1-5 The image anomaly detection method based on a pre-trained diffusion model described in this invention includes the following steps:

[0036] Step 1: From the publicly available ImageNet dataset, select 30 classes for image anomaly detection. Choose one class as the normal class and the remaining 29 classes as the anomaly class. During network training, only the normal class images are used as the training dataset, while all 30 classes are used as the test dataset. All images are uniformly cropped to 256×256 pixels. Strong and weak data augmentation are then performed on the cropped training dataset to obtain augmented images. Then, an enhanced image training dataset {D} is constructed. The test dataset is simply cropped to a size of 256×256 without any other data augmentation operations.

[0037] Strong data augmentation S refers to data augmentation methods that produce significant changes to an image. Strong data augmentation constructs an image that has certain similarities to the original image but does not conform to the original image's distribution patterns. This invention includes four types of strong data augmentation: rotating the image clockwise by 0°, 90°, 180°, and 270°, denoted as S1, S2, S3, and S4 respectively. The expression for strong data augmentation S is as follows:

[0038] S={S1,S2,S3,S4} (1)

[0039] Weak data augmentation (W) refers to a data augmentation method that produces only subtle changes to an image. It constructs an image that differs slightly from the original but still exhibits a similar distribution pattern. Weak data augmentation employs a random augmentation approach, meaning it performs weak data augmentation multiple times on the same image, each time yielding a different result, thereby constructing more images that conform to a normal image distribution. In this invention, weak data augmentation includes: randomly cropping and enlarging the image to 256×256, randomly changing the image's grayscale, and randomly changing the image's color.

[0040] For an input image x, after undergoing the i-th type of strong data augmentation S... i Enhanced image obtained after weak data augmentation W The expression is as follows:

[0041]

[0042] For an input dataset {B}, where x∈{B}, the size of the dataset is |B|. All input images x need to be processed by all strong data augmentation methods to generate different strongly augmented images, resulting in a number of k×B strongly augmented images. Then, weak data augmentation is applied twice to each strongly augmented image, ultimately resulting in a number of 2k×B augmented images. Let {D} represent the entire augmented image training dataset.

[0043] Proceed to step 2.

[0044] Step 2: Construct a global image anomaly detection network, such as... Figure 1 As shown, it includes a text-to-image diffusion model, an implicit description generator, and a compression module. This will enhance the image. The implicit description generator is used to obtain the text embedding e, which is then used as a condition to enhance the image. A common input image diffusion model is used to obtain the semantic representation r of the image. The semantic representation is then compressed to obtain image features. The specific steps are as follows:

[0045] Step 2.1: Obtain the semantic representation of the image using a text-to-image diffusion model, such as... Figure 2 As shown, a stable diffusion model is used as the text-to-image diffusion model in this invention. The stable diffusion model is trained on a large text-image dataset, and its network structure includes a text encoder, an autoencoder, and a U-Net network. The text encoder transforms the input text description into a text embedding e, and stable diffusion controls the image generation process based on this text embedding e. Stable diffusion uses an autoencoder to map the image to a latent space, and then uses a U-Net network in the latent space to predict the noise of the input image using the text embedding e and time step t. The text embedding of stable diffusion influences the intermediate output of the U-Net network through a cross-attention mechanism, making the final generated image closely related to the input text. Therefore, we believe that the intermediate output of the U-Net network contains rich semantic information, and that appropriate text embeddings can influence the content of semantic information. This invention uses the U-Net network, and at time step t=0, uses the text embedding e as a condition to obtain... The semantic representation of text. Text embedding e is introduced in step 2.2.

[0046] The UNet network can be further divided into input modules, intermediate modules, and output modules. Each module contains several intermediate blocks, which can be convolutional blocks, residual blocks, or attention blocks. The k-th intermediate output layer M of the UNet network's output module... k Its shape is represented as |D|×c k ×h k ×w k c k h is the number of output channels for the k-th layer. k w represents the height of the output of the k-th layer. k M represents the width of the output at the k-th layer. k The expression is as follows:

[0047]

[0048] Where t represents the time step, and in this invention t = 0. e represents text embedding. Select k This indicates that the kth intermediate block of the UNet output module is selected for output.

[0049] Further processing is performed on the output of each layer, h k ×w k Upsampled to the same size h×w and then stitched together, this is used as the semantic representation r of the image. The expression for this process is shown in equation (4):

[0050] r=Concatenate([M0,M3,M6…,M n (4)

[0051] Where n represents the total number of selected intermediate blocks, this invention selects the concatenation of the outputs of every 3 output blocks in the output module as the semantic representation. Concatenate represents the concatenation operation. The final shape of the semantic representation is |D|×c×h×w, where c is the sum of the number of output channels of each layer, as shown in equation (5):

[0052]

[0053] Step 2.2: Construct an implicit description generator, such as... Figure 3 As shown, the stable diffusion model can generate high-quality images based on the input text description. The text description is transformed into a text embedding e by a text encoder, and then injected into the UNet network through an attention mechanism. UNet is greatly influenced by the text embedding e during the computation process, and a suitable text embedding e can enhance the semantic representation of the image. This invention proposes an implicit description generator to generate the corresponding text embedding e of an image without text input. Compared with text description, this method has a wider range of applications and more accurate descriptions. The implicit description generator consists of a Contrastive Language-Image Pretraining (CLIP) model and linear layers. The CLIP model is also the text encoder used in the stable diffusion model, which can map images and text to a similar embedding vector space. This invention utilizes the CLIP model to enhance the image. The vector space is mapped to an embedding vector space that is close to the text space, and then transformed into the text embedding space through a linear layer. The expression for this process is shown in Equation (6):

[0054]

[0055] Step 2.3: Construct the compression module. Figure 4 The structure of the compression module is shown, which includes convolutional blocks, average pooling layers (AvgPool), and linear layers. Each convolutional block consists of a 3×3 convolutional layer, a batch normalization layer (BatchNorm), and a ReLU activation function. The compression module compresses the semantic representation r to obtain image features. In addition, the compression module includes an extra strong augmentation classification layer to determine the strong classification type used for the image. This layer consists of a linear layer to enhance the robustness of the model. Proceed to step 3.

[0056] Step 3: Construct the loss function, which includes feature contrast loss and strong enhancement classification loss. The specific steps are as follows:

[0057] Step 3.1: Constructing the Feature Contrast Loss. To ensure the image features remain discriminative, this invention introduces a contrastive learning method to train the network. The goal of contrastive learning is to make similar images closer together in the feature space, while dissimilar images are kept as far apart as possible. Through contrastive learning, the network learns more discriminative feature representations of the images. In contrastive learning, similar images are called positive pairs, and dissimilar images are called negative pairs. The feature contrast loss function L... sim The expression is shown in equation (7):

[0058]

[0059] In this loss function, for any augmented image in the augmented image training dataset {D} {X +}express Positive examples for the dataset, It is a set {X + Any member of}. In this invention, multiple enhanced images obtained by performing multiple weak enhancements on the same strongly enhanced image are considered as positive pairs. -}express Negative examples of the dataset, It is a set {X - Any member in}, in this invention, The negative pairs are all the augmented images except for the positive pairs. Specifically, as shown below... Figure 5 As shown. The `sim` function is used to calculate the similarity between two features. In this invention, `sim` represents the cosine similarity, that is, for two feature values ​​a and b, the cosine similarity is `sim(a,b) = (a·b) / ||a|| ||b||`, where a higher value indicates a higher similarity, meaning the two features are closer in the feature space. `z(·)` represents the output of the feature contrast layer, and `τ` is the temperature of the feature contrast loss, which determines the weight of the contrast loss function on positive and negative pairs. In this invention, its value is taken as 0.5.

[0060] Step 3.3: Construct the strongly enhanced classification loss L cls For the strong enhancement classification layer, the cross-entropy loss function is used to calculate the loss. Its inputs are the output CLS(·) of the strong enhancement classification layer and the i-th strong enhancement type actually used for the enhanced image. The specific expression is as follows:

[0061]

[0062] Step 3.4: The total loss function L is shown in (10), where α is the total loss function of L. cls The weight.

[0063] L = L sim +α×L cls (10)

[0064] Proceed to step 4.

[0065] Step 4: Train the entire anomaly detection network using the enhanced image training dataset.

[0066] The anomaly detection network is trained using an enhanced image training dataset. Throughout the process, the parameters of the CLIP model and the UNet network are kept constant; only the parameters of the linear layers and compression modules of the implicit description generator need to be trained and updated to obtain the trained anomaly detection network. Proceed to step 5.

[0067] Step 5: Input the test dataset into the trained anomaly detection network, calculate the similarity between the test dataset and the features of all normal class images in the enhanced image training dataset, and select the value with the highest similarity as the anomaly score of the image. For the test image y in the test dataset, its score(y) expression is as shown in (11):

[0068]

[0069] This invention uses an anomaly score to represent the degree of anomaly of a test image; the lower the score, the more likely the test image is to be abnormal. The feature similarity between similar images in the network model is also higher.

Claims

1. An image anomaly detection method based on a text-to-image diffusion model, characterized in that, Includes the following steps: Step 1: From the publicly available ImageNet dataset, select 30 classes for image anomaly detection; select one class as the normal class and the remaining 29 classes as the anomaly class; during network training, only the normal class images are used as the training dataset, and during testing, all 30 classes of images are used as the test dataset; crop all images to a uniform size of 256×256, and perform strong data augmentation and weak data augmentation on the cropped training dataset to obtain augmented images. This leads to the construction of an enhanced image training dataset. ; The test dataset is simply cropped to a size of 256×256 without any other data augmentation operations, and then proceeds to step 2. Step 2: Construct a global image anomaly detection network: The overall image anomaly detection network includes a text-to-image diffusion model, an implicit description generator, and a compression module; it enhances the image... Input implicit description generator to obtain text embedding Then, text embedding is used as a condition, along with image enhancement. A common input image diffusion model is used to obtain the semantic representation of the image. The semantic representation is used to obtain image features through the compression module; Constructing a global image anomaly detection network: The overall image anomaly detection network includes a text-to-image diffusion model, an implicit description generator, and a compression module; it enhances the image... Input implicit description generator to obtain text embedding Then, text embedding is used as a condition, along with image enhancement. A common input image diffusion model is used to obtain the semantic representation of the image. The semantic representation is compressed to obtain image features. The specific steps are as follows: Step 2.1: Obtain semantic representations of images using a text-to-image diffusion model. A stable diffusion model was adopted as the text-to-image diffusion model. The stable diffusion model was trained on a large text-image dataset, and its network structure included a text encoder, an autoencoder, and a U-Net network. The text encoder transformed the input text description into text embeddings. stable diffusion depends on the text embedding Control the image generation process; Stable diffusion uses an autoencoder to map an image to a latent space, and then uses a UNet network on the latent space for text embedding. and time step Predict noise in the input image; Stable diffusion text embedding influences the intermediate output of the UNet network through a cross-attention mechanism, making the final generated image closely related to the input text. Using the UNet network, at time step... In the case of text embedding As a condition, obtain semantic representation; Step 2.2: Construct an implicit description generator. The stable diffusion model can generate high-quality images based on the input text description. The text description is converted into text embeddings by a text encoder. Then, it is injected into the UNet network through an attention mechanism; the implicit description generator consists of a contrastive language image pre-trained CLIP model and linear layers, which uses the CLIP model to enhance the image. The vector space is mapped to an embedding vector space that is close to the text space, and then transformed back to the text embedding space through a linear layer. The expression for this process is as follows: (6) Step 2.3: Construct a compression module, which includes convolutional blocks, average pooling layers, and linear layers. Each convolutional block consists of a 3×3 convolutional layer, a batch normalization layer, and a ReLU activation function. The compression module compresses semantic representations. Compression is performed to obtain image features; in addition, the compression module also includes an extra strong enhancement classification layer consisting of a linear layer to determine the strong classification type used by the image, thereby enhancing the robustness of the model. Proceed to step 3; Step 3: Construct the loss function, which includes feature contrast loss and strong enhancement classification loss. The specific steps are as follows: Step 3.1: Construct the feature contrast loss: A contrastive learning method is introduced to train the network. In contrastive learning, similar images are called positive pairs, and dissimilar images are called negative pairs. The feature contrastive loss function is used. The expression is as follows: （7） In feature contrast loss function In the context of augmenting image training datasets Any enhanced image , express Positive examples for the dataset, It is a set Any member in; multiple enhanced images obtained by performing multiple weak enhancements on the same strongly enhanced image are considered as positive pairs; express Negative examples of the dataset, It is a set Any member in, The function is used to calculate the similarity between two features. This is represented as the output of the feature contrast layer. The temperature of the feature contrast loss determines the weight of the contrast loss function on positive and negative pairs. Step 3.3: Construct a strongly enhanced classification loss. : For the strong augmentation classification layer, the cross-entropy loss function is used to calculate the loss, and its input is the output of the strong augmentation classification layer. And the first method actually used to enhance the image Strongly enhanced typing, the specific expression is as follows: （9） Step 3.4, Total Loss Function as follows: （10） in for The weights; Proceed to step 4; Step 4: Train the entire anomaly detection network using the augmented image training dataset: The anomaly detection network is trained using the enhanced image training dataset. Throughout the process, the parameters of the CLIP model and the UNet network are fixed. Only the parameters of the linear layer and the compression module of the implicit description generator need to be trained and updated to obtain the trained anomaly detection network. Proceed to step 5. Step 5: Input the test dataset into the trained anomaly detection network, calculate the similarity of features between the test dataset and all normal class images in the augmented image training dataset, and select the highest similarity value as the anomaly score for the image. For the test images in the test dataset... Its score as follows: Anomaly scores are used to represent the degree of anomaly in test images. The lower the score, the more likely the test image is to be abnormal, and the higher the feature similarity between similar images in the network model.

2. The image anomaly detection method based on the text-to-image diffusion model according to claim 1, characterized in that, In step 1, for the input image After going through the first Seed strength data enhancement and weak data augmentation The resulting enhanced image The expression is as follows: （2） For the input dataset , Its quantity is All input images It is necessary to use all the strong data augmentation methods to generate different strongly augmented images, at which point the number of strongly augmented images is... Then, weak data augmentation is applied twice to each strongly augmented image, resulting in a final number of augmented images. ,use This represents the entire augmented image training dataset. .

3. The image anomaly detection method based on the text-to-image diffusion model according to claim 2, characterized in that, Set up strong data enhancement There are four types, including rotating the image clockwise by 0°, 90°, 180°, and 270°, denoted as follows: Strong data enhancement The expression is as follows: （1） Weak data augmentation This refers to data augmentation methods that only make subtle changes to the image, creating an image that has subtle differences from the original image but has similar distribution patterns through weak data augmentation; Weak data augmentation adopts a random augmentation method, that is, weak data augmentation is performed on the same image multiple times, and the results are different each time, so as to construct more images that conform to the normal image distribution; Weak data augmentation This includes randomly cropping and enlarging the image to 256×256, randomly changing the grayscale of the image, and randomly changing the color of the image.

4. The image anomaly detection method based on the text-to-image diffusion model according to claim 1, characterized in that, In step 2.1, the UNet network is divided into an input module, intermediate modules, and an output module. Each module contains several intermediate blocks, and the types of intermediate blocks include convolutional blocks, residual blocks, and attention blocks. The output module of the UNet network... Intermediate output layer Its shape is represented as , For the first The number of output channels of the layer. For the first The height of the layer output, For the first The width of the layer output; The expression is as follows: （3） in Represents the time step, let ; Indicates text embedding; Indicates the selection of the output module of UNet. One intermediate block output; Process the output of each layer, Upsampled to the same size The images are then stitched together to form a semantic representation of the image. The expression for this process is as follows: （4） in This represents the total number of selected intermediate blocks. This represents a concatenation operation, and the final semantic representation is as follows: ,in The sum of the output channels for each layer is expressed as follows: （5）。

Images