Image compressed sensing reconstruction method based on generation prior diffusion

Abstract

The invention belongs to the technical field of signal processing, and particularly relates to an image compressed sensing reconstruction method based on generation prior diffusion. According to the method, the problem of insufficient utilization of prior information in an existing data-driven image compressed sensing method is solved, a deep expansion network and a diffusion generation model are fused, and a two-section image reconstruction method based on generation of prior diffusion is provided. In the training of the first stage, a variation encoder is used for learning compact prior expression of a natural image, and a dynamic near-end mapping module is designed in a deep expansion reconstruction network to embed prior information. And in the second stage, freezing parameters of the prior learning model, inferring image prior by applying the lightweight diffusion model, and embedding the image prior into the deep expansion network to complete image reconstruction. Experimental results on a plurality of reference sets show that compared with the existing method, the image compressed sensing reconstruction method provided by the invention can improve the reconstruction quality of the natural image.

Description

Technical Field

[0001] The invention discloses an image compression sensing reconstruction method based on generating a priori diffusion, and belongs to the technical field of signal processing. Background Art

[0002] Compressed Sensing (CS) is a new information acquisition paradigm, the core of which is to reconstruct signals from very few measurements with the help of structural priors. This feature inspires natural image acquisition in energy-constrained systems, greatly reducing the requirements for sensor performance and communication bandwidth, and alleviating the burden of data storage and transmission. The breadth of CS has brought many impressive applications, including but not limited to single-pixel imaging, accelerated magnetic resonance imaging, snapshot compressed imaging, wireless sensor networks and remote sensing. The application of CS to natural image acquisition mainly faces three core problems: sampling strategy, reconstruction method and prior setting.

[0003] Formally, the image perception process can be expressed as: y=Φx in is the original image, is the measurement matrix, is the measured value. The goal of the reconstruction algorithm is to find an approximate solution for x using the acquired y and the given Φ. Obviously, this is an underdetermined inverse problem. To obtain an accurate estimate, it is necessary to fully exploit the prior information of natural images and impose restrictions on Φ. Early studies have shown that Gaussian random matrices are a suitable choice for CS. However, it gives the same weight to all components, which may not be optimal for image acquisition. Previous work has shown that learned matrices help improve reconstruction performance, and the present invention also uses this strategy for sampling.

[0004] Image reconstruction can be formulated as solving the following optimization problem: in. is a regularization term used to introduce prior knowledge. λ is a weight coefficient. Early work introduced a variety of iterative algorithms to solve this problem, including Approximate Message Passing (AMP), Proximal Gradient Descent (PGD), and Alternating Direction Method of Multipliers (ADMM). PGD is suitable for solving many large-scale linear inverse problems. The image reconstruction problem is solved by iterating between the following two update steps: where k = 1, 2, ..., K is the iteration stage, λ is the gradient descent step size, Φ T is the transpose of the measurement matrix.

[0005] As neural networks demonstrate powerful feature learning capabilities, many data-driven CS methods have emerged. These methods can be roughly divided into two categories: black-box models and deep unfolding networks (DUN). Black-box models attempt to directly establish an inverse mapping from measurement space to image space, and optimize network parameters through end-to-end training on image datasets. For example, CSNet uses the residual structure as the backbone network and improves the image reconstruction quality by jointly optimizing the sampling module and the reconstruction module; CSformer uses a hybrid framework of Convolutional Neural Networks (CNN) and Transformer to simultaneously capture the local features and global associations of the image. The black-box model can implicitly summarize the feature information of natural images, has a simple structure, and is highly efficient. However, this type of method bypasses the key point of image priors, and the reconstruction process is difficult to explain and lacks mathematical guarantees.

[0006] DUN attempts to use neural networks to simulate traditional iterative methods. For example, ISTA-Net uses CNN-based nonlinear units to simulate the sparse transformation process and deeply expands the iterative threshold shrinkage algorithm (ISTA) to complete iterative reconstruction; AMP-Net uses stacked convolutional layers to replace the denoising module in AMP to improve the quality of reconstructed images. TransCS hybridizes CNN and Transformer to simultaneously capture local features and global associations of images. DPC-DUN uses a dynamic path controllable depth expansion network to adaptively select a fast and appropriate reconstruction path for different images. The above work has laid a good foundation for the development of interpretable deep CS methods, but existing methods focus on the improvement of sampling and reconstruction models, and there is still a lack of research on deep priors.

[0007] At present, some works have explored the pattern of deep priors in image CS. Among them, generative priors are an effective strategy. Such methods attempt to solve the inverse problem in a compact low-dimensional space and perform well at extremely low sampling rates. However, image dimensionality reduction will inevitably lead to loss of high-frequency information. In addition, generative models usually require that the distribution of training samples is concentrated, so the reconstruction method based on generation has poor effect on image reconstruction outside the range of the generator and is difficult to generalize to a wide range of natural images. In recent years, diffusion models (DM) have shined in the field of image generation and density estimation. Many works have applied DM to the underlying visual field and achieved satisfactory results.

[0008] DM models the unknown data distribution p(z) by learning to invert a Markov chain of length T. Its forward diffusion process gradually destroys the original data structure and generates a series of noisy data {z1,z2,…,z T}, this process can be formally expressed as: where z t is the noisy data at time step t, β t ∈(0,1) is an increasing variance table, i.e., β t-1 <β t The entire diffusion process is parameter-free and converges approximately to a standard Gaussian distribution after T steps of iteration. The reverse denoising process of the conditional DM learns a parameterized inverse process to infer the original data, which is also Markov-like and can be formally expressed as: where Ι is a unit vector, c is an embedding condition guiding the generated direction, and μ θ (z t ,c,t) is used to predict the posterior mean: where ∈ θ (·) is the neural network for predicting superimposed noise.

[0009] Inspired by the above work, this paper integrates DUN and DM and proposes a deep expansion framework based on generative prior diffusion for image CS. Summary of the invention

[0010] The purpose of the present invention is to embed data-driven generative prior into DUN and propose an image CS reconstruction method based on generative prior diffusion.

[0011] This approach is a two-stage architecture. Figure 1As shown in the figure. In the first stage, the natural image prior applicable to image CS is summarized by jointly training the Generative Prior Learning Module (GPLM) and the CNN-based DUN to obtain a compact expression of natural images in a low-dimensional space. After the training is completed, a compact prior space can be obtained to embed the iterative reconstruction process. In the second stage, the parameters of the generative prior learning module are frozen, and the generative prior diffusion module (GPDM) and DUN are jointly trained. In GPDM, the initial reconstruction is encoded as a condition to guide the generation process to infer the compact expression of the image to be tested. Finally, the decoded image compact prior is embedded in the dynamic proximal mapping module (DPMM) in the iterative reconstruction process to improve the quality of the reconstructed image.

[0012] Assuming that the original image to be observed in the CS process is x, the sampling module first uses the operation Split x into non-overlapping blocks of size C×B×B. Then use the operation (·) Expand the image block to n = CB 2 dimensional vector. Finally, the learned measurement matrix is ​​used Get measurement values Formally, the process of the sampling module can be described as:

[0013] The initial reconstruction is the inverse process of the sampling module. First, the learned initial reconstruction matrix is ​​used The measurement y is mapped to an n-dimensional vector and then the operation υ- 1 (·) Reshape it into a C×B×B image block and finally use the operation Stitched into a complete image, overall, the initial reconstruction module is described as:

[0014] The structure of GPLM is shown in the attached Figure 2 , is a variant of the variational autoencoder based on CNN, including an encoder and a decoder, which are optimized only in the first stage of training. In the encoder, a 3×3 convolutional layer is used to extract the feature map of the original image, and then the size of the feature map is gradually reduced and the number of channels is increased. Downsampling is achieved by average pooling with a kernel size of 2×2. Residual blocks (RB) are used to process feature maps at each scale, and their structure is similar to Figure 1The RB in is the same. Finally, the encoder uses a 3×3 convolutional layer to output the mean μ and variance σ of the compact representation of the image. Formally, this process can be expressed as: μ,σ=ε(x)

[0015] The decoder is the reverse process of the encoder, using It first samples Then the feature map size is gradually increased and the number of channels is reduced, and finally the decoded prior is output by a 3×3 convolutional layer, where upsampling is achieved by nearest neighbor interpolation.

[0016] The structure of GPDM is shown in the attached Figure 3 , including two steps, diffusion and generation, which are optimized only in the second stage of training. GPDM uses the one-stage pre-trained GPLM to obtain the compact prior distribution of the original image x, and uses its mean as the original data, that is, z0 = μ. The forward diffusion process is non-learning, and noise is gradually added to z0 in the following way: in After T steps of iteration, it converges approximately to the standard Gaussian distribution. The reverse generation process is also Markov-like, with z T As input, the added noise is gradually removed using the iterative process shown in the following formula: where ∈ θ (·) is a parameterized neural network consisting of a 3×3 convolution input, an output layer, and 5 RBs in the middle. After T steps of iteration, the image compact representation is inferred.

[0017] The iterative reconstruction module deeply expands the PGD algorithm and uses K stages to reconstruct the original image. The reconstruction block structure of each stage is shown in the attached Figure 1 , including the Gradient Descent Module (GDM) and DPMM. The input of the reconstruction block is divided into three parts, namely the reconstruction of the previous stage And the feature map h k-1 , and the compact expression z0 (first stage) or (Second stage). h0 is obtained by applying a single layer of convolution to x0. GDM is processed by the following formula get where λ k is the gradient descent step size for learning. Next, concatenate With h k-1 Get the input of DPMM The overall processing flow of DPMM can be described as follows: in It is the backbone of DPMM, consisting of the input layer, output layer and 4 RBs in the middle. The parameters of are optimized in the first stage of training and frozen in the second stage. The compact expression is provided by GPLM in the first stage of training and inferred by GPDM in the second stage of training. Finally, after K stages of iteration, the final reconstructed image is obtained.

[0018] In the first stage of training, the sampling module, reconstruction module and GPLM are jointly optimized. The training loss is defined as follows: Where ρ is the weight parameter, which is a very small constant. is the Kullback-Leibler divergence, which is used to constrain the distribution of image compact priors. In the second stage of training, the parameters of GPLM are frozen, and the sampling module, reconstruction module and GPDM are jointly optimized. The training loss is defined as follows: Where γ is the weight parameter. BRIEF DESCRIPTION OF THE DRAWINGS

[0019] Figure 1 This is the overall architecture of an image compression sensing reconstruction method based on generating a priori diffusion proposed by the present invention.

[0020] Figure 2 It is the network structure of the generative prior learning module designed by the present invention.

[0021] Figure 3 It is a generation prior diffusion module framework designed by the present invention.

[0022] Figure 4 This is a comparison of the reconstructed samples between the present invention and some existing methods when the sampling rate is 0.1. DETAILED DESCRIPTION

[0023] In order to more clearly illustrate the implementation scheme and significance and advantages of the present invention, the present invention is described in more detail below in conjunction with the accompanying drawings and implementation examples.

[0024] (1) Training set preparation: The present invention produces a two-stage training set based on BSD500. Specifically, each image in the dataset is randomly cropped into 400 image blocks of size 64×64, and a variety of methods are used, including random scaling, random rotation, horizontal flipping, and vertical flipping, to improve data diversity. The present invention uses the BSD500 dataset to maintain consistency with the comparison algorithm, and a larger training set can further improve the performance of this method.

[0025] (2) Sampling module: The original image x is cut into image blocks of size 32×32 and expanded into vectors. Each image block contains n=32×32=1024-dimensional inputs. Then, the learned measurement matrix Get the m-dimensional measurement value y. Assuming the sampling rate is 0.1, the dimension of each image block measurement value is m = 0.1 × 1024 = 102. During training, initialize the measurement matrix Φ as a Gaussian random matrix.

[0026] (3) Initial reconstruction module: using the learned reconstruction matrix Project the m-dimensional measurements of each image patch into an n-dimensional vector, reshape it and concatenate it into the initial reconstruction During training, the reconstruction matrix A is initialized to the transpose of Φ.

[0027] (4) GPLM: See the attached figure for the overall structure Figure 2 , including two parts: encoder and decoder. The encoder uses 3×3 convolution to extract the feature map of a single-channel grayscale image. The feature map size is 64×64 and the number of channels is 32. Then two RBs are used to process the feature map. The first RB outputs 32 channels and the second RB outputs 64 channels. Its jump connection uses 1×1 convolution to adjust the number of channels. Then 2×2 average pooling is used to reduce the feature map size to 32×32. The processing on the second size is similar to the first size. The final feature map size is 16×16 and the number of channels is 128. At the smallest size, one RB is first used to process the feature map, and then a 3×3 convolution is used to obtain a 2-channel output, where the first channel is the mean μ of the compact expression and the second channel is the variance σ of the compact expression. The structure of the decoder is symmetrical with that of the encoder. From the distribution The sampling is compactly expressed in z0, and a 128-channel feature map is extracted using a 3×3 convolution. Then, RB is used to process the feature map and reduce the number of channels, and the nearest neighbor interpolation is used to increase the feature map size. Finally, a 3×3 convolution is used to output a 64×64 single-channel decoded feature map.

[0028] (5) GPDM: See the attached for the overall structure Figure 3, including two processes: diffusion and generation. The diffusion process uses the pre-trained GPLM encoder to obtain a compact expression of the original image, and uses the mean of its distribution as the original data z0=μ; then T steps are used to iterate and add noise to z0 to obtain a variable z close to Gaussian noise. T The generation process is based on z T Input, with the initial reconstruction encoding As a condition, a denoising block composed of a neural network is used to iteratively remove the noise added in z T The noise on S2 (·) The structure is exactly the same as the encoder in GPLM except for the output layer, which only outputs the condition of a single channel. Denoising block splicing z t , and time step t as input, use 3×3 convolution to extract 32-channel feature maps, then use 5 RBs to process the feature maps, and finally use 3×3 convolution to output the denoised variable z t-1 After T steps of iteration, the inferred compact expression is obtained

[0029] (6) Iterative reconstruction module: Use K stages to reconstruct the original image. The reconstruction block structure of each stage is shown in the attached Figure 1 , including GDM and DPMM. GDM performs gradient descent operation, except for the step size λ k It is not learning. After executing gradient descent, we get the intermediate variables Splicing With the feature map h of 63 channels k-1 Get 64-channel input of DPMM where h0 is a 3×3 convolution operation on DPMM first splices and A 65-channel feature map is obtained; then the number of channels is adjusted to 64 using 3×3 convolution, and 4 consecutive RBs and 1 layer of 3×3 convolution are used to process the feature map; finally, the processed feature map is compared with Add and cut to get the reconstruction of the current stage And the feature map h k .

[0030] (7) Compact expression distribution learned from GPLM in the first stage The sample z0 is decoded and sent to the reconstruction block, and the sampling, initial reconstruction, iterative reconstruction module and GPLM are jointly optimized. The loss function is defined as follows: Where ρ is the weight parameter, which is a very small constant. is the Kullback-Leibler divergence, which is used to constrain the distribution of image compact priors. The second stage is to compactly express the GPDM inference The decoded data is sent to the reconstruction block, and the sampling, initial reconstruction, iterative reconstruction module and GPDM are jointly optimized. The loss function is defined as follows: Where γ is the weight parameter.

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

[0032] The simulation experiment software of the present invention is configured as follows: Ubuntu 20.04 operating system, the simulation language is Python, and the software library is Pytorch1.7.0.

[0033] In the simulation experiment, three widely used natural image sets, Set14, BSD100 and Urban100, were used to test and compare the quality of the reconstructed images of the present invention and the existing methods. The comparative experiment compares the image compression sensing reconstruction method based on generative prior diffusion proposed by the present invention with the existing advanced algorithms, including CSNet+, AMP-Net, TransCS, DPC-DUN and CSformer. Among them, CSNet+ and CSformer are black box models, and AMP-Net, TransCS and DPC-DUN are DUN models.

[0034] The comparison indicators are Peak Signal-to-Noise Ratio (PSNR) and Structural Similarity (SSIM). The higher the PSNR and SSIM, the better the reconstruction effect.

[0035] In the experiment, the sampling rates were set to 0.04, 0.10, 0.25 and 0.5 respectively. Table 1 shows the average PSNR and SSIM of the reconstructed images of the present invention and the existing advanced methods on different data sets. It can be seen that the present invention is superior to the comparison method in all data sets and all sampling rates in the table. Specifically, the average PSNR / SSIM of the reconstructed images of the present invention in 3 data sets (214 images) and 5 sampling rates is 30.50 / 0.8623, which is 0.77 / 0.0269 (2.59% / 3.22%) higher than the second best DUN method (TransCS), and 0.79 / 0.0160 (2.65% / 1.89%) higher than the best black box model (CSformer).

[0036] Attached Figure 4 The reconstruction effect of the ppt3 image in the Set14 dataset by the present invention and the comparison method is shown when the sampling rate is 0.1. From the local detail enlargement image, it can be seen that the present invention can reconstruct the texture of the image more accurately. Table 1 PSNR / SSIM comparison between the present invention and existing advanced methods at different sampling rates

Claims

1. A compressed sensing image reconstruction method based on generative prior diffusion, comprising five submodules: sampling, initial reconstruction, generative prior learning, generative prior diffusion and iterative reconstruction, characterized in that: (1) The sampling module first uses the operation Cut the original image x into non-overlapping blocks of C×B×B, and then use the operation Expand the image block to n = CB 2 dimensional vector, and finally use the learned measurement matrix Get measurement values Overall, the submodule is described as (2) Initial reconstruction is the inverse process of the sampling module. First, the learned initial reconstruction matrix is ​​used Map the measurement y to an n-dimensional vector and then use the operation Reshape it into a C×B×B image block and finally use the operation Stitched into a complete image, Overall, the submodule is described as (3) The generative prior learning module consists of two parts: the encoder and the decoder. The encoder consists of a convolutional layer, a residual block, and a downsampling operation. The output mean and variance μ, σ = (x), compactly express z0 from the distribution The decoder is the inverse process of the encoder, consisting of a convolutional layer, a residual block, and an upsampling operation, and outputs a single-channel feature map (4) Generate a priori The diffusion module includes diffusion and generation processes. Forward diffusion gradually adds noise to z0, and finally obtains a Gaussian noise variable z T , reversely generate T As input, through iteration Inferring a compact representation of an image Denoise(·) consists of a convolutional layer and a residual block. S2 (·) encodes the initial reconstruction as a condition suitable for the diffusion model, and has the same structure as (·) except for the output layer. t=1,2,,T is the time step; (5) The iterative reconstruction module consists of K stages, where the input of stage k includes the compact representation z0 and the reconstruction of the previous stage And the feature map h k-1 , the reconstruction block first performs gradient descent to obtain Then splice With h k-1 Get the intermediate feature map Finally, the dynamic proximal mapping operation is performed to output the reconstruction and feature map of this stage. in It consists of 2 convolutional layers and 4 residual blocks; (6) The training of the image reconstruction model is divided into two stages. In the first stage, the compact expression z0 obtained by the generative prior learning module is decoded and input into the iterative reconstruction module, which is jointly optimized with the sampling and initial reconstruction modules to learn the image generative prior suitable for compressed sensing. In the second stage, the parameters of the generative prior learning module are frozen, and the generative prior diffusion module is trained with the learned compact prior z0 to decode the image compact prior inferred by it. The iterative reconstruction module pre-trained in the first stage is input, and the image reconstruction is completed by joint optimization with the pre-trained sampling and initial reconstruction modules.

Images