A natural image compressive sensing method based on latent diffusion model

Abstract

The present application belongs to the technical field of signal processing, and particularly relates to a natural image compressive sensing method based on a latent diffusion model. The present application deeply explores the discreteness and compressibility of natural images, and proposes a two-stage image reconstruction method based on deep learning. In the first stage, the vector quantization generative adversarial network is used to learn the deep features of the image on a large data set to obtain the discretized encoding of the image in the low-dimensional latent space, so as to eliminate the redundant information of the natural image. In the second stage, the diffusion model is used to infer the corresponding latent encoding of the image under the condition of the compressed measurement value of the image, and the original image is further reconstructed. The experimental results show that, compared with the existing method, the image compressive sensing method proposed by the present application can greatly improve the visual effect of the reconstructed image under a low sampling rate.

Description

[Technical Field]

[0001] This invention discloses a natural image compressed sensing method based on a latent diffusion model, belonging to the field of signal processing technology. [Background Technology]

[0002] Images are important carriers of information, giving rise to numerous applications in the era of big data. Traditional image acquisition systems typically first sample all pixels in the field of view to form a raw image, and then use complex compression algorithms to remove redundant information from the image for storage and transmission. This requires image sensors with strong computing power and sufficient energy, which may not be suitable for some energy-constrained and inexpensive sensor applications, such as large-scale wireless Internet of Things (IoT).

[0003] Compressed sensing (CS) breaks the limitations of the Nyquist-Shannon sampling theorem, enabling signal reconstruction from a small number of measurements. Furthermore, by merging sampling and compression processes, it directly obtains measured values ​​using linear projection, reducing the computational demands and energy consumption of sensors, making it a highly efficient image acquisition method. Today, image CS is widely used in various fields, such as MRI, single-pixel cameras, fast imaging, and holographic imaging.

[0004] The CS process can be described as follows:

[0005]

[0006] in For the original image, For measured values, It is a measurement matrix CS aims to utilize a given and Solving the underdetermined linear system shown in the above formula yields the original image. An approximate solution.

[0007] From a probabilistic perspective, the goal of the reconstruction algorithm is to maximize the posterior probability. To determine the variable According to Bayesian theory:

[0008]

[0009] Where for a given , It is a constant. It is the likelihood function, determined by the observation model, and is generally obtained by minimizing the fidelity term. Please provide a solution. Prior knowledge derived from images can be solved by constructing regularization terms. Common image priors include sparsity in the transform domain, low rank, and nonlocal self-similarity. Bayesian-based algorithms solve the CS problem by embedding prior knowledge into the likelihood inference process. The main research challenge is the prior setting problem.

[0010] In recent years, researchers have proposed several deep learning-based algorithms to solve the CS reconstruction problem. Compared with traditional algorithms, these algorithms can automatically learn prior information about images through neural networks and directly apply this information through training on datasets. Modeling offers higher accuracy and faster speed. Furthermore, in deep CS algorithms, the sampling matrix is ​​learned rather than hand-designed, making it more suitable for specific applications. These algorithms can be broadly categorized into two types: data-driven algorithms and depth unrolling algorithms. The former learns image priors by stacking convolutional layers, acting as a black-box model. The latter aims to design suitable networks to replace certain modules in model-based methods, such as sparse transforms and denoisers.

[0011] Purely data-driven algorithms cannot explain the reconstruction process, and their convolutional layers focus on local image features, ignoring global information. Depth unrolling algorithms still essentially utilize known image priors. Finding more suitable prior information remains an open problem. Most existing image CS reconstruction methods aim to improve peak signal-to-noise ratio (PSNR), but their evaluation results may differ significantly from human perception. A major drawback is that these metrics can lead to overly smooth reconstructed images. How to improve the perceptual quality of reconstructed images remains an open problem.

[0012] Meanwhile, diffusion models have demonstrated powerful capabilities in image generation. By defining a Markov chain of expansion steps, they transform unknown data distributions into a known distribution, such as a standard Gaussian distribution. This is equivalent to... Modeling. The Latent Diffusion Model (LDM) introduces an autoencoder to generate images using a diffusion model in a low-dimensional latent space. This removes redundant information from the images, improves computational efficiency, and is more suitable for likelihood-based generative models. [Summary of the Invention]

[0013] The purpose of this invention is to apply the latent diffusion model to the compressed sensing reconstruction process of natural images, and to propose a compressed sensing method for natural images based on the latent diffusion model.

[0014] The objective of this invention is to measure values By setting conditions to guide the generation process of LDM, high-quality reconstructed images were obtained. .

[0015] Assume the original image to be observed during the CS process is Its potential encoding is The observed value is The objective of this invention is to train a parameterized reconstruction network to approximate... According to Bayes' theorem, we can obtain:

[0016]

[0017] For a given measurement value, Since the measured value is constant, it is only related to the original image and the sampling matrix. Therefore, the reconstruction process is decomposed into the following three sub-problems:

[0018] (1) Maximize : Is for It is sufficient, that is This invention solves the likelihood estimation problem by designing a parameterized sampling module;

[0019] (2) Maximize This problem is modeled by a pre-trained vector quantization generative adversarial network decoder.

[0020] (3) Maximize This is the key point of the present invention, which is solved by training conditional LDM.

[0021] The purpose of this invention is to solve the problem of depth compressed sensing reconstruction of natural images in energy-constrained systems through the following technical solution:

[0022] This invention uses a learnable sampling matrix based on block compressed sensing to establish a linear mapping from image space to measurement space. Specifically, the natural image is first divided into segments of size [missing information]. Given a non-overlapping block with 3 channels, give... One input. Then use the operation. Expand each block as 3D vectors, and use learnable measurement matrices. Sampling was performed to obtain Dimensional measurement values. For example, for When the sampling rate is 0.1, the sampling module will provide... One output. The sampling module is described as... :

[0023]

[0024] The reconstruction module of this invention includes three sub-modules: conditional encoding, latent variable reconstruction, and decoding.

[0025] In LDM, the dimensions of the conditions are always adjusted to be the same as the latent variables in order to perform the splicing operation, even though the number of channels may differ. This invention uses a learnable initial reconstruction matrix and a downsampling network to encode measurements into conditions usable in LDM. First, the initial reconstruction matrix is ​​used... Map the measured values ​​to the initial reconstructed vector; then use the operation... The vector is reshaped into an image; finally, a downsampling network is used to encode the initial reconstructed image into the conditions required for the LDM. Specifically, assume the image size is... latent variables The size is .for The sampling block, measurement The size is First, use the initial reconstruction matrix. Map the measured values ​​to dimensions. The reconstructed vector. Then use the operation. Reshape the reconstructed vector to size. The initial image. Finally, a convolutional neural network-based downsampling network is used. Encode the image into a size of conditions The conditional coding submodule is described as follows: :

[0026]

[0027] The LDM module is designed to be based on The diffusion process is used to predict the encoding of the original image. Assume the original image... The potential encoding is The forward diffusion process involves continuously moving towards... Adding noise to obtain the sequence For any step ,have:

[0028]

[0029] in , ,and The latent variable reconstruction module can be described as a conditional inverse diffusion process:

[0030]

[0031] Specifically, this invention predicts the first... The noise added step by step, and through The next iteration of denoising yields... The focus of reconstruction is predicting noise at each step. This invention uses a UNet with cross-attention to model a conditional temporal denoising autoencoder. In the steps splicing and The splicing serves as the input to UNet in the current step, therefore [ z t , y c ] The shape is UNet obtains feature maps of different sizes through four downsampling operations and preserves details by skipping connections within the same size. After sufficient training, the reconstruction module... and Given the input, predict the latent encoding of the original image:

[0032]

[0033] This invention uses a pre-trained decoder Reconstruct the original image. Vector quantization generative adversarial networks discretize and encode the latent features of the image, which is more consistent with the modalities found in nature. The decoder will first... Each element in the code is replaced with the closest code in the code table. Then a decoder is used to reconstruct the code.

[0034]

[0035] This invention is an end-to-end CS structure capable of measuring... The input is the reconstructed image. During training, the sampling and reconstruction modules are jointly optimized. The set of all parameters is... The training objective is to minimize the denoised score matching loss.

[0036] L( Oh )= E x,ϵ ,t [ ∥ ϵ- ϵ I ( z t ,t, C C,Q ( S A ( x ))) ∥ 2 2 ]

[0037] [Advantages and positive effects of the present invention]

[0038] Compared with the prior art, the present invention has the following advantages and positive effects:

[0039] First, this invention designs a data-driven sampling paradigm and corresponding encoding module for perceptually enhanced CS. Unlike traditional methods that use manually designed sampling matrices, the sampling matrix in this invention is jointly optimized with the reconstruction network. Furthermore, this invention designs an initialization encoding network that encodes measurements into conditions to guide the latent diffusion model in reconstructing the low-dimensional latent encoding of the original image.

[0040] Second, this invention designs a latent diffusion-based CS reconstruction model and explains its working principle. It shifts the image reconstruction work from the complex image space to a low-dimensional latent space, and uses a conditional diffusion model to simultaneously solve the likelihood estimation and data prior problems in the Bayesian CS problem, thus reconstructing a low-dimensional representation of the image;

[0041] Third, this invention proposes an end-to-end two-stage image CS architecture. This architecture uses a pre-trained autoencoder to establish a mapping from the image space to a low-dimensional latent space, and uses a diffusion model to obtain latent vectors conditioned on the measurement values, thereby reconstructing the image. Experiments show that it can reconstruct images with high quality at low sampling rates while preserving key details. [Attached Image Description]

[0042] Figure 1 This is a flowchart of a deep compressed sensing image reconstruction method for wireless sensor networks proposed in this invention;

[0043] Figure 2 This is a comparison of the reconstruction examples of natural images by the present invention and existing advanced algorithms at a sampling rate of 0.04;

Detailed Implementation Methods

[0044] To make the embodiments and advantages of the present invention clearer, the present invention will be described in more detail below with reference to the accompanying drawings and examples.

[0045] (1) The sampling module first samples the original color natural image Cut into dimensions of Image blocks, then each image block contains The system takes one input and unfolds each image patch into a vector, then uses a learnable measurement matrix. ,in ,Will 3D input projection is Dimensional measurement value For transmission and storage, the matrix is ​​measured during training. It is initialized as a Gaussian random matrix;

[0046] (2) The initial encoding module first uses a learnable reconstruction matrix , measure value Projection as The 1D vector is then reshaped and stitched together to form the initial reconstructed image. Then, a learnable downsampling network is used to encode the initial reconstruction into conditions that can be input to the diffusion module. The downsampling network uses two The average pooling yields a compression ratio of 4 and a channel count of 5. At each size, two convolutional layers with the "ReLU" activation function are used to process the feature map, and the feature map has 64 channels.

[0047] (3) The latent diffusion module includes a forward process and a backward process. The forward process first uses pre-trained vector quantization to generate the encoder of the adversarial network to obtain the latent encoding of the original image. Then, noise of increasing intensity is added to it, resulting in a series of noisy codes. },in Approximately standard Gaussian noise, the backward process uses random noise. and conditions As input, UNet progressively predicts the noise added to the image encoding, ultimately recovering the original image. Latent encoding The UNet has 4 layers, and the number of channels in the feature map of each layer is 160, 320, 640 and 1280 respectively. An attention mechanism is used on the smallest feature map.

[0048] (4) Use gradient descent to iteratively optimize steps (1)-(3) to obtain the optimal sampling matrix. Reconstruction Matrix Given the initial encoding network and UNet parameters, the loss function is:

[0049] L( Oh )= E x,ϵ ,t [ ∥ ϵ- ϵ I ( z t ,t, C C,Q ( S A ( x ))) ∥ 2 2 ]

[0050] (5) After training, Denoising Diffusion Implicit Models (DDIM) are used to accelerate the inference process of LDM to measure values. The input is used to obtain the predictive encoding. ;

[0051] (6) The decoding module quantizes the pre-trained vectors to generate the decoder of the adversarial network, using the encoding predicted by the latent diffusion module. For input, first target For each element, find the closest encoding in the pre-trained codebook and replace it with that encoding. The corresponding elements are then processed through multiple convolutions and upsampling to obtain the reconstructed image. The decoder consists of two upsampling networks with a scaling factor of 2 and an output layer. The first upsampling network contains two residual blocks with 128 channels and a bilinear interpolation operation. The second upsampling network contains two residual blocks with 64 channels and a bilinear interpolation operation. The output layer is a single-layer convolutional neural network with 64 input channels and 3 output channels.

[0052] The hardware configuration for the simulation experiment of this invention is as follows: Intel(R) Core(R) i5-13600KF CPU; 32.0GB DDR4 memory; NVIDIA Quadro RTX3090 GPU.

[0053] The simulation software configuration of this invention is as follows: Ubuntu 20.04 operating system, Python simulation language, and PyTorch 1.7.0 software library.

[0054] In the simulation experiment, a training set was constructed using BSD500. All images were randomly cropped into 60 images of size [missing information]. The subgraphs used data augmentation techniques include random rotation, horizontal folding, and vertical folding. Therefore, the training set includes 10,000 subgraphs of size [missing information]. Images were used for testing. Five widely used natural image datasets were used: Set14, BSD100, Urban100, and DIV2K. Comparative experiments compared the proposed natural image compressed sensing method based on a latent diffusion model with existing state-of-the-art algorithms, including ISTA-Net, OPINE-Net, AMP-Net, TransCS, and TCS-Net. The comparison metric was Learned Perceptual ImagePatch Similarity (LPIPS). LPIPS is a commonly used evaluation metric in low-level vision, and compared to Peak Signal-to-Noise Ratio (PSNR) and Structural Similarity (SSIM), it better aligns with human perception. Lower LPIPS indicates better reconstruction results.

[0055] Table 1. Comparison of LPIPS between the present invention and existing advanced methods at different sampling rates.

[0056]

[0057] In the experiments, sampling rates were set to 0.01, 0.04, and 0.1, respectively. Table 1 shows the average LPIPS of the reconstructed images using the present invention and existing state-of-the-art methods on different datasets. It can be seen that the present invention outperforms the comparison methods on all datasets and at all sampling rates in the table. Furthermore, the advantages of the present invention become more pronounced as the sampling rate decreases. For example, on the Urban100 dataset, when the sampling rate is 0.1, the LPIPS of the present invention is 0.0815, which is 0.0393 lower than the second-best TransCS; when the sampling rate decreases to 0.01, the LPIPS of the present invention is 0.1874, which is 0.1465 lower than the second-best TCSNet. Therefore, the present invention is more suitable for image compression sensing tasks at low sampling rates and has better application prospects.

[0058] Figure 2 Examples of Lenna image reconstruction using the present invention and the comparison method at a sampling rate of 0.04 are shown. It can be seen that the image reconstructed by the present invention is clearer, contains richer details, and has a better visual effect.

Claims

1. A natural image compressed sensing method based on a latent diffusion model, comprising four sub-modules: a sampling module, an initial encoding module, a latent diffusion module, and a decoding module, characterized in that: (1) The sampling module first samples the original color natural image Cut into dimensions of Image blocks, then each image block contains The system takes one input and unfolds each image patch into a vector, then uses a learnable measurement matrix. ,in ,Will 3D input projection is Dimensional measurement value For transmission and storage, the matrix is ​​measured during training. It is initialized as a Gaussian random matrix; (2) The initial encoding module first uses a learnable reconstruction matrix , measure value Projection as The 1D vector is then reshaped and stitched together to form the initial reconstructed image. Then, a learnable downsampling network is used to encode the initial reconstruction into conditions that can be input to the diffusion module. The downsampling network uses two The average pooling layer is used to obtain a compression ratio of 4 and a channel number of 5. Two convolutional layers with "ReLU" activation function are used to process the feature map at each size, and the feature map has 64 channels. (3) The latent diffusion module includes a forward process and a backward process. The forward process first uses pre-trained vector quantization to generate the encoder of the adversarial network to obtain the latent encoding of the original image. Then, noise of increasing intensity is added to it, resulting in a series of noisy codes. },in Approximately standard Gaussian noise, the backward process uses random noise. and conditions As input, noise is progressively predicted and added to the image encoding using UNet, and the final output is the original image. Latent encoding The UNet has 4 layers, and the number of channels in the feature map of each layer is 160, 320, 640 and 1280 respectively. An attention mechanism is used on the smallest feature map. (4) The decoding module quantizes the pre-trained vectors to generate the decoder of the adversarial network, using the encoding predicted by the latent diffusion module. For input, first target For each element, find the closest encoding in the pre-trained codebook and replace it with that encoding. The corresponding elements are then processed through multiple convolutions and upsampling to obtain the reconstructed image. The decoder consists of two upsampling networks with a scaling factor of 2 and an output layer. The first upsampling network contains two residual blocks with 128 channels and a bilinear interpolation operation. The second upsampling network contains two residual blocks with 64 channels and a bilinear interpolation operation. The output layer is a single-layer convolutional neural network with 64 input channels and 3 output channels.

Images