An improved DDPM-based contraband image generation method
By improving the multi-scale feature extraction and feature fusion of the DDPM model, the problem of insufficient quality and diversity of contraband image generation caused by the instability of GANs training was solved, and more efficient contraband image generation was achieved.
Patent Information
- Authority / Receiving Office
- CN · China
- Patent Type
- Patents(China)
- Current Assignee / Owner
- CIVIL AVIATION UNIV OF CHINA
- Filing Date
- 2023-10-27
- Publication Date
- 2026-07-03
Smart Images
Figure CN117409099B_ABST
Abstract
Description
Technical Field
[0001] This invention belongs to the field of computer vision and image processing technology, and specifically relates to a method for generating images of contraband based on an improved DDPM (denoising diffusion probability model). Background Technology
[0002] Image generation of contraband is an important task in the field of computer vision, used to generate synthetic images that resemble real images of contraband, often for expanding datasets. Traditional image generation methods are often based on Generative Adversarial Networks (GANs), but due to problems such as unstable training and pattern collapse, the quality and diversity of the generated images are limited. Summary of the Invention
[0003] To address the aforementioned problems, the present invention aims to provide a method for generating images of contraband based on an improved DDPM.
[0004] To achieve the above objectives, the method for generating images of contraband based on the improved DDPM provided by the present invention includes the following steps performed in sequence:
[0005] Step 1: Create a dataset of prohibited items;
[0006] Step 2: Construct a contraband image generation model based on the improved DDPM and input the contraband dataset obtained in Step 1 above to generate feature maps of contraband in different poses;
[0007] Step 3: Use the contraband dataset created in Step 1 to train the contraband image generation model based on the improved DDPM constructed in Step 2. Combine the feature maps of contraband in different poses generated in Step 2 to determine the weights of the contraband image generation model based on the improved DDPM.
[0008] Step 4: Generate an image of a certain contraband using the improved DDPM-based contraband image generation model with the weights trained in Step 3.
[0009] In step 1, the specific method for creating the contraband dataset is as follows:
[0010] First, multiple raw X-ray images of a certain contraband item in different poses are acquired using a security X-ray machine. Then, the contraband item in the raw X-ray images is labeled using the image annotation tool Labelme to obtain a JSON file. Next, the coordinate positions in the JSON file are used to convert the raw X-ray images into a mask image. Then, the mask image is subjected to dilation and erosion calculations to obtain a Trimap image. The Trimap image is then converted into an Alpha image, and the Alpha image is overlaid with the raw X-ray image to extract the foreground, thus obtaining preliminary sample images. The preliminary sample images are then proportionally reduced to a size of 64×64 to obtain sample images M. All sample images M constitute the contraband dataset Q.
[0011] In step 1, the prohibited items include lighters, matches, liquid containers, and power banks.
[0012] In step 2, the specific method for constructing a contraband image generation model based on the improved DDPM and inputting it into the contraband dataset obtained in step 1 above to generate feature maps of contraband in different poses is as follows:
[0013] 2.1 First, construct a contraband image generation model based on the improved DDPM, and then input the contraband dataset Q obtained in step 1 into the above-mentioned contraband image generation model based on the improved DDPM;
[0014] 2.2 Secondly, each sample image M in the input contraband dataset Q is encoded. The encoding process includes three downsampling operations, as follows: The first downsampling involves taking each sample image M after passing through two 1×1 convolutional layers and performing three max pooling operations using 3x3, 5x5, and 7x7 pooling kernels with a stride of 2 when passing through the mixed banded pooling layer MSP, resulting in three feature maps M1, M2, and M3 at different scales. Then, the obtained feature maps M1, M2, and M3 are processed along the channel dimension. The features are concatenated to obtain a feature map M4 with multi-scale information. Finally, feature map M4 is passed through a sequence consisting of a 1x1 convolutional layer, a depthwise separable convolutional layer, and another 1x1 convolutional layer to obtain a feature map P with 128 channels, which serves as the result of the first downsampling. The second downsampling involves passing the feature map P obtained from the first downsampling after two 1x1 convolutional layers through a hybrid banded pooling layer (MSP) and performing three max pooling operations using 3x3, 5x5, and 7x7 pooling kernels with strides of 2, respectively, to obtain the third feature map P. We obtain feature maps P1, P2, and P3 at different scales. Then, we concatenate the obtained feature maps P1, P2, and P3 along the channel dimension to obtain a feature map P4 with multi-scale information. Finally, we pass feature map P4 through a sequence consisting of a 1x1 convolutional layer, a depthwise separable convolutional layer, and a 1x1 convolutional layer to obtain a feature map F with 256 channels, which serves as the result of the second downsampling. The third downsampling is performed by taking the feature map F obtained from the second downsampling after passing through two 1×1 convolutional layers and then passing it through hybrid band pooling. During the MSP layer, three max pooling operations are performed using pooling kernels with strides of 2 (3x3, 5x5, and 7x7) to obtain three feature maps F1, F2, and F3 at different scales. Then, the feature maps F1, F2, and F3 are concatenated along the channel dimension to obtain a feature map F4 with multi-scale information. Finally, feature map F4 is passed through a sequence consisting of a 1x1 convolutional layer, a depthwise separable convolutional layer, and a 1x1 convolutional layer to obtain a feature map X with 512 channels, which serves as the result of the third downsampling.
[0015] 2.3 Then, the decoding operation is performed. The decoding process is achieved through multiple repeated convolution and upsampling operations. The specific process is as follows: The feature map X obtained by the third downsampling after two convolutional layers is used as the input for the first upsampling. Then, it is upsampled through a 2x2 convolutional layer to obtain a feature map X1 with 512 channels after fusing with the skip connection features. The feature map X1 is then passed through the AC-ECA module to expand the effective receptive field of the convolution kernel. This module consists of dilated convolutional layers AC and ECA attention mechanisms. Then, the feature map X1 is passed through a 1x1 convolutional layer to obtain feature map X2. The feature map X2 is used as the input for the second upsampling. The feature map X3, with 256 channels, is obtained by upsampling through a 2x2 convolutional layer and fusing it with the skip connection features. X3 is then passed through an AC-ECA module to expand the effective receptive field of the convolutional kernel. Finally, X3 is passed through a 1x1 convolutional layer to obtain feature map X4. X4 is then used as input for a third upsampling operation, passing through another 2x2 convolutional layer to obtain feature map X5, with 128 channels, fusing it with the skip connection features. X5 is then passed through an AC-ECA module to expand the effective receptive field of the convolutional kernel. Finally, X5 is passed through a 1x1 convolutional layer to obtain the final output feature map X6.
[0016] 2.4 Finally, the pixels in the feature map X6 are classified using the Softmax activation function and a 1×1 convolutional layer, and a feature map A representing the prediction result is output with the same size as the input sample image M.
[0017] In step 3, the contraband image generation model based on the improved DDPM, constructed in step 2, is trained using the contraband dataset created in step 1. The specific method for determining the weights of the contraband image generation model based on the improved DDPM, by combining the feature maps of contraband in different poses generated in step 2, is as follows:
[0018] 3.1 Input the contraband dataset Q obtained in step 1 as the data source into the above-mentioned contraband image generation model based on the improved DDPM;
[0019] 3.2 Set the initial weights for model training and begin iterative training;
[0020] 3.3. Compare the feature map A representing the prediction result obtained in step 2 with the corresponding sample image M in the contraband dataset Q and calculate the loss function. Then, use the backpropagation algorithm to adjust the weights of the model training to obtain the weights after this round of training.
[0021] 3.4 After reaching the required number of iterations, output the weights of the contraband image generation model based on the improved DDPM.
[0022] In step 4, the specific method for generating an image of a certain contraband using the contraband image generation model based on the improved DDPM with the weights trained in step 3 is as follows:
[0023] Using the weights obtained from training in step 3, a contraband image generation model based on the improved DDPM is used to output a generated image of a certain contraband.
[0024] The method for generating images of contraband based on the improved DDPM provided by this invention has the following beneficial effects:
[0025] (1) The present invention changes the original DDPM from a five-layer structure to a four-layer structure, thereby reducing the number of model parameters and improving the running speed of the model.
[0026] (2) The present invention replaces the max pooling in the original DDPM downsampling with the MSP-DSC module. After the max pooling operation at different scales, a depth-separable convolutional layer is introduced. It uses max pooling operation at different scales to downsample the input image. The purpose is to increase the model’s ability to perceive features at different scales, which can improve the model’s performance, perception ability and efficiency, and can better handle multi-scale features, while reducing the model’s parameters and computational load.
[0027] (3) In this invention, the first convolutional layer after upsampling in each feature fusion process is replaced with an AC-ECA module, which consists of dilated convolutional layers and an ECA attention mechanism. Dilated convolution expands the effective receptive field of the convolutional kernel by introducing holes in the kernel. By increasing the receptive field, the model can better capture global contextual information in the image, thereby improving the accuracy of image segmentation. The ECA attention mechanism can dynamically adjust the weights of different regions according to the content of the image, making the model pay more attention to important regions in the image. Attached Figure Description
[0028] Figure 1 This is a flowchart of the method for generating images of contraband based on the improved DDPM provided by the present invention;
[0029] Figure 2 This is a schematic diagram of a contraband image generated using the improved DDPM-based contraband image generation method provided by this invention. Detailed Implementation
[0030] The present invention will now be described in detail with reference to the accompanying drawings and specific embodiments.
[0031] like Figure 1 As shown, the method for generating images of contraband based on the improved DDPM provided by this invention includes the following steps performed in sequence:
[0032] Step 1: Create a dataset of prohibited items;
[0033] First, multiple raw X-ray images of a certain contraband item in different poses are acquired using a security X-ray machine. Then, the contraband item in the raw X-ray images is labeled using the image annotation tool Labelme to obtain a JSON (JavaScript Object Notation) file. Next, the coordinate positions in the JSON file are used to convert the raw X-ray images into a mask image. Then, the mask image is subjected to dilation and erosion calculations to obtain a Trimap image. The Trimap image is then converted into an Alpha image and overlaid with the raw X-ray image to extract the foreground, obtaining a preliminary sample image. The preliminary sample image is then proportionally reduced to a size of 64×64 to obtain a sample image M. All sample images M constitute the contraband dataset Q.
[0034] The prohibited items include lighters, matches, liquid containers, and power banks; the prohibited items dataset created by this invention contains a total of 3,800 sample images.
[0035] Step 2: Construct a contraband image generation model based on the improved DDPM and input the contraband dataset obtained in Step 1 above to generate feature maps of contraband in different poses;
[0036] This invention uses a contraband image generation model based on an improved DDPM to generate feature maps of contraband in different poses. Its structure is a typical encoder-decoder framework, and the main process includes four parts: image input, encoding operation, decoding operation, and image output.
[0037] The specific method is as follows:
[0038] 2.1 First, construct a contraband image generation model based on the improved DDPM, and then input the contraband dataset Q obtained in step 1 into the above-mentioned contraband image generation model based on the improved DDPM;
[0039] 2.2 Next, each sample image M in the input contraband dataset Q is encoded. The encoding process is the feature extraction process, which includes three downsampling operations. The downsampling is implemented by the MSP-DSC module in the contraband image generation model based on the improved DDPM. This module includes a mixed-stripe pooling layer (MSP) and a depthwise separable convolution layer (DSC). The MSP-DSC module uses max pooling operations at different scales to downsample the input sample image M, with the aim of increasing the model's ability to perceive features at different scales. The specific process is as follows: The first downsampling involves taking each sample image M after passing through two 1×1 convolutional layers and performing triple max pooling with 3x3, 5x5, and 7x7 kernels with a stride of 2 when passing through a hybrid banded pooling layer (MSP). This yields three feature maps M1, M2, and M3 at different scales. Then, the feature maps M1, M2, and M3 are concatenated along the channel dimension to obtain a feature map M4 with multi-scale information. Finally, feature map M4 is passed through a depthwise separable convolutional layer consisting of a 1x1 convolutional layer. The sequence of layers and 1x1 convolutional layers yields a feature map P with 128 channels, which serves as the result of the first downsampling. The second downsampling involves passing the feature map P obtained from the first downsampling after two 1x1 convolutional layers through a hybrid banded pooling layer (MSP) with three max pooling operations using 3x3, 5x5, and 7x7 kernels with strides of 2, resulting in three feature maps P1, P2, and P3 at different scales. These feature maps P1, P2, and P3 are then concatenated along the channel dimension to obtain a feature map with... The feature map P4 contains multi-scale information. Finally, feature map P4 is passed through a sequence consisting of a 1x1 convolutional layer, a depthwise separable convolutional layer, and another 1x1 convolutional layer to obtain a feature map F with 256 channels, which serves as the result of the second downsampling. The third downsampling involves passing the feature map F obtained from the second downsampling after two 1x1 convolutional layers through a hybrid banded pooling layer (MSP) using three max pooling operations with 3x3, 5x5, and 7x7 kernels with strides of 2, respectively, to obtain feature maps of three different scales. The model first performs downsampling on features F1, F2, and F3. Then, it concatenates these features along the channel dimension to obtain feature map F4, which has multi-scale information. Finally, feature map F4 is passed through a sequence consisting of a 1x1 convolutional layer, a depthwise separable convolutional layer, and another 1x1 convolutional layer to obtain feature map X with 512 channels, which serves as the result of the third downsampling. Through these three downsampling steps, the model gradually reduces the size of the feature map while increasing the number of channels, thus achieving multi-scale extraction of image features.
[0040] 2.3 Next, the decoding operation is performed. The decoding process is achieved through multiple repeated convolution and upsampling operations. Specifically, the feature map X obtained from the third downsampling after two convolutional layers is used as the input for the first upsampling. Then, it is upsampled through a 2x2 convolutional layer to obtain a feature map X1 with 512 channels after fusing with the skip connection features. Feature map X1 is then passed through the AC-ECA module to expand the effective receptive field of the convolutional kernel. This module consists of atrous convolution layers (AC) and an Efficient Channel Acceleration (ECA) attention mechanism. The model constructs an attention mechanism, then passes feature map X1 through a 1x1 convolutional layer to obtain feature map X2. Feature map X2 is then used as input for a second upsampling, followed by a 2x2 convolutional layer to obtain feature map X3 with 256 channels after fusing with skip connection features. Feature map X3 is then passed through an AC-ECA module to expand the effective receptive field of the convolutional kernel, followed by a 1x1 convolutional layer to obtain feature map X4. Feature map X4 is then used as input for a third upsampling, followed by a 2x2 convolutional layer to obtain feature map X5 with 128 channels after fusing with skip connection features. Feature map X5 is then passed through an AC-ECA module to expand the effective receptive field of the convolutional kernel, followed by a 1x1 convolutional layer to obtain the final output feature map X6. Through these three upsampling operations, the model gradually restores the size of the feature maps while reducing the number of channels, achieving layer-by-layer fusion and recovery of image features.
[0041] 2.4 Finally, the pixels in the feature map X6 are classified using the Softmax activation function and a 1×1 convolutional layer, and a feature map A representing the prediction result is output with the same size as the input sample image M.
[0042] Step 3: Use the contraband dataset created in Step 1 to train the contraband image generation model based on the improved DDPM constructed in Step 2. Combine the feature maps of contraband in different poses generated in Step 2 to determine the weights of the contraband image generation model based on the improved DDPM.
[0043] Setting up the training environment: Train the improved model using an HP deep learning workstation running Ubuntu 16.04, with the CPU version specified. The processor is 4110 with a processing frequency of 2.20GHz; the graphics card is an NVIDIA GeForce GTX 2080Ti; it is equipped with CUDA 11.4 and CUDNN 8.2 to accelerate the computation of the network model; the Python version is 3.7; the deep learning framework used for training is PyTorch 1.7.0; the number of training iterations is 1100; the batch size is set to 32; the number of convolution channels is set to 64; the input image size is 64×64; the Adam optimizer is used; the maximum learning rate is set to 0.01; the minimum learning rate is set to 0.0001; and the learning rate decay is achieved using cosine annealing.
[0044] The specific method is as follows:
[0045] 3.1 Input the contraband dataset Q obtained in step 1 as the data source into the above-mentioned contraband image generation model based on the improved DDPM;
[0046] 3.2 Set the initial weights for model training and begin iterative training;
[0047] 3.3. Compare the feature map A representing the prediction result obtained in step 2 with the corresponding sample image M in the contraband dataset Q and calculate the loss function. Then, use the backpropagation algorithm to adjust the weights of the model training to obtain the weights after this round of training.
[0048] 3.4 After reaching the required number of iterations, output the weights of the contraband image generation model based on the improved DDPM; the number of iterations in this invention is 1100.
[0049] Step 4: Generate an image of a certain contraband using the improved DDPM-based contraband image generation model with the weights trained in Step 3.
[0050] Using the weights trained in step 3, the improved DDPM-based contraband image generation model outputs a generated image of a certain contraband, such as... Figure 2 As shown.
[0051] To verify the effectiveness of this invention, the inventors tested the contraband image generation model based on the improved DDPM and evaluated the quality of the generated images using the FID (Firmware Identifier) index, which assesses the quality of the generation model. A lower FID value indicates a smaller difference between the generated image and the real image, and thus a higher quality generation model. The results are shown in Table 1.
[0052] Table 1
[0053]
[0054]
[0055] As shown in Table 1, the method of the present invention performs better than the original model in generating images of contraband, and the difference between the generated images and the real images is smaller, thus achieving the goal of expanding the dataset.
Claims
1. A method for generating images of contraband based on an improved DDPM, characterized in that: The method for generating images of contraband includes the following steps performed in sequence: Step 1: Create a dataset of prohibited items; Step 2: Construct a contraband image generation model based on the improved DDPM and input the contraband dataset obtained in Step 1 above to generate feature maps of contraband in different poses; The specific method is as follows: 2.1 First, construct a contraband image generation model based on the improved DDPM, and then input the contraband dataset Q obtained in step 1 into the above-mentioned contraband image generation model based on the improved DDPM; 2.2 Secondly, each sample image M in the input contraband dataset Q is encoded. The encoding process includes three downsampling operations, as follows: The first downsampling involves taking each sample image M after passing through two 1×1 convolutional layers and performing three max pooling operations using 3x3, 5x5, and 7x7 pooling kernels with a stride of 2 when passing through the mixed banded pooling layer MSP, resulting in three feature maps M1, M2, and M3 at different scales. Then, the feature maps M1, M2, and M3 are concatenated along the channel dimension to obtain a feature map M4 with multi-scale information. Finally, feature map M4 is... A sequence of 1x1 convolutional layers, depthwise separable convolutional layers, and 1x1 convolutional layers is used to obtain a feature map P with 128 channels, which serves as the result of the first downsampling. The second downsampling involves passing the feature map P obtained from the first downsampling after two 1x1 convolutional layers through a hybrid banded pooling layer (MSP) using three max pooling operations with 3x3, 5x5, and 7x7 kernels with strides of 2, resulting in three feature maps P1, P2, and P3 at different scales. These feature maps P1, P2, and P3 are then concatenated along the channel dimension to obtain a feature map P4 with multi-scale information. Finally, feature map P4... The second downsampling process involves passing the feature map F, which has 256 channels, through a sequence of 1x1 convolutional layers, depthwise separable convolutional layers, and 1x1 convolutional layers. The third downsampling process involves passing the feature map F, obtained after the second downsampling process through two 1x1 convolutional layers, through a mixed banded pooling layer (MSP) with three max pooling operations using 3x3, 5x5, and 7x7 kernels with a stride of 2, resulting in three feature maps F1, F2, and F3 at different scales. These feature maps F1, F2, and F3 are then concatenated along the channel dimension to obtain a feature map F4 with multi-scale information. Finally, feature map F4 is passed through a sequence of 1x1 convolutional layers, depthwise separable convolutional layers, and 1x1 convolutional layers to obtain a feature map X, which has 512 channels, and this is the result of the third downsampling process. 2.3 Then, the decoding operation is performed. The decoding process is achieved through multiple repeated convolution and upsampling operations. The specific process is as follows: The feature map X obtained by the third downsampling after two convolutional layers is used as the input for the first upsampling. Then, it is upsampled through a 2x2 convolutional layer to obtain a feature map X1 with 512 channels after fusing with the skip connection features. The feature map X1 is then passed through the AC-ECA module to expand the effective receptive field of the convolution kernel. This module consists of dilated convolutional layers AC and ECA attention mechanisms. Then, the feature map X1 is passed through a 1x1 convolutional layer to obtain feature map X2. The feature map X2 is used as the input for the second upsampling. The feature map X3, with 256 channels, is obtained by upsampling through a 2x2 convolutional layer and fusing it with the skip connection features. X3 is then passed through an AC-ECA module to expand the effective receptive field of the convolutional kernel. Finally, X3 is passed through a 1x1 convolutional layer to obtain feature map X4. X4 is then used as input for a third upsampling operation, passing through another 2x2 convolutional layer to obtain feature map X5, with 128 channels, fusing it with the skip connection features. X5 is then passed through an AC-ECA module to expand the effective receptive field of the convolutional kernel. Finally, X5 is passed through a 1x1 convolutional layer to obtain the final output feature map X6. 2.4 Finally, the softmax activation function and a 1×1 convolutional layer are used to classify each pixel in the feature map X6 above, and a feature map A representing the prediction result is output with the same size as the input sample image M. Step 3: Use the contraband dataset created in Step 1 to train the contraband image generation model based on the improved DDPM constructed in Step 2. Combine the feature maps of contraband in different poses generated in Step 2 to determine the weights of the contraband image generation model based on the improved DDPM. Step 4: Generate an image of a certain contraband using the improved DDPM-based contraband image generation model with the weights trained in Step 3.
2. The method for generating images of contraband based on improved DDPM according to claim 1, characterized in that: In step 1, the specific method for creating the contraband dataset is as follows: First, multiple raw X-ray images of a certain contraband item in different poses are acquired using a security X-ray machine. Then, the contraband item in the raw X-ray images is labeled using the image annotation tool Labelme to obtain a JSON file. Next, the coordinate positions in the JSON file are used to convert the raw X-ray images into a mask image. Then, the mask image is subjected to dilation and erosion calculations to obtain a Trimap image. The Trimap image is then converted into an Alpha image. The Alpha image is then overlaid with the raw X-ray image to extract the foreground and obtain a preliminary sample image. Then the initial sample image is scaled down to 64. Sample images M are obtained with a size of 64, and the contraband dataset Q is composed of all sample images M.
3. The method for generating images of contraband based on improved DDPM according to claim 1, characterized in that: In step 1, the prohibited items include lighters, matches, liquid containers, and power banks.
4. The method for generating images of contraband based on improved DDPM according to claim 1, characterized in that: In step 3, the contraband image generation model based on the improved DDPM, constructed in step 2, is trained using the contraband dataset created in step 1. The specific method for determining the weights of the contraband image generation model based on the improved DDPM, by combining the feature maps of contraband in different poses generated in step 2, is as follows: 3.1 Input the contraband dataset Q obtained in step 1 as the data source into the above-mentioned contraband image generation model based on the improved DDPM; 3.2 Set the initial weights for model training and begin iterative training; 3.
3. Compare the feature map A representing the prediction result obtained in step 2 with the corresponding sample image M in the contraband dataset Q and calculate the loss function. Then, use the backpropagation algorithm to adjust the weights of the model training to obtain the weights after this round of training. 3.4 After reaching the required number of iterations, output the weights of the contraband image generation model based on the improved DDPM.
5. The method for generating images of contraband based on improved DDPM according to claim 1, characterized in that: In step 4, the specific method for generating an image of a certain contraband using the contraband image generation model based on the improved DDPM with the weights trained in step 3 is as follows: Using the weights obtained from training in step 3, a contraband image generation model based on the improved DDPM is used to output a generated image of a certain contraband.