NL-CS Net: Deep Learning with Non-Local Pior for Image Compressive Sensing

Compressed sensing (CS) theory has received a lot of attention in recent years. CS proves that when a signal is sparse in a certain domain, it can be recovered with high probability from much fewer measurements than the Nyquist sampling theorem ref1 ; ref2 ; ref3 ; ref4 ; ref5 . The potential reduction in measurements is attractive for diverse practical applications, including but not limited to magnetic resonance imaging (MRI) ref6 , radar imaging ref7 and sensor networks ref8 .

Over the past decades, a great deal of image CS reconstruction methods have been developed based on sparse representation model ref9 , which operates on the assumption that many images can be sparsely represented by a dictionary. The majority of those traditional methods use some structured sparsity as an image prior and then solve a sparsity-regularized optimization problem in an iterative fashion ref10 ; ref11 . Some elaborate structures were introduced into CS, like Gaussian scale mixtures model in wavelet domain ref12 . In addition, non-local self-similarity of image is also introduced to enhance the CS performance ref13 ; ref14 ; ref15 . For example, Metzler $et$ $al$. ref16 combined a Block Matching 3D (BM3D) denoiser into approximate message passing (AMP) framework to perform CS reconstruction. Zhang $et$ $al$. ref13 proposed the method combining sparse prior with non-local regularizers achieved well performance. Recently, some optimization-based methods have implemented adaptive sampling using alternating optimization techniques to jointly optimize the sampling matrix and CS recovery algorithms ref9 . Despite the excellent interpretability of the above methods, they all require hundreds of iterations to produce decent results, which inevitably entails a heavy computational burden, in addition to the challenges posed by hand-craft transformations and associated hyper-parameters.

Inspired by the successful applications of deep learning, several network-based CS reconstruction methods were developed to learn the inverse mapping from the CS measurement domain to original signal domain ref17 ; ref18 ; ref19 . Mousavi $et$ $al$. ref20 applied a stacked denoising auto-encoder (SDA) to learn the statistical relationship from training data. However, the fully connected network used in SDA results in high computation cost. Kulkarni $et$ $al$. ref21 developed a method based on convolutional neural networks, called Recon-Net, to reconstruct the original image from the CS sampled image blocks. Yao $et$ $al$. ref22 used residual learning to further improve CS reconstruction. Sun $et$ $al$. ref23 propose a novel sub-pixel convolutional generative adversarial network (GAN) to learn compressed sensing reconstruction of images. To mitigate block effect in reconstruction, some models make use of full image areas for reconstruction ref24 ; ref25 . Meanwhile, for further improving the CS performance, some models are proposed to train the non-linear recovery operator to learn the optimal sampling pattern with recovery model ref26 ; ref27 ; ref28 . The main advantage of the network-based methods is the reconstruction speed, as opposed to their optimization-based counterparts. However, the barrier to future performance improvement is their lack of the CS domain-specific insights intrinsic to optimization-based approaches.

To overcome above shortcomings, researchers link optimization methods to networks, which make them interpretable. Specifically, these methods embed the solving process of traditional optimization-based methods into the forward operator of deep learning. For instance, Zhang $et$ $al.$ ref29 proposed a deep network called ISTA-Net, which maps the popular Iterative Shrinkage Thresholding (ISTA) algorithm to network. It learns the sparse transform and soft threshold function via network. Based on ISTA-Net, Zhang $et$ $al$. ref30 proposes Opine-Net, which combines an efficient sampling module with ISTA-Net to achieve adaptive sampling and recovery. More recently, You $et$ $al$. ref31 improved ISTA-Net, which enables a single model to adapt for multiple sampling rates. Xiang $et$ $al$. ref32 proposed FISTA-Net for solving inverse problem, which is an accelerated version of ISTA. The Alternating Direction Method of Multipliers (ADMM) is proposed for the saddle point problem containing Lagrange multipliers that can not be solved directly by the ISTA algorithm. Drawing on the same idea, Yang $et$ $al$. ref33 proposed the ADMM-Net, which unfolds ADMM into network and applies it to CS-MRI. It employs a learnable transformation and the corresponding hyper-parameters in ADMM are learned from the network. Zhang $et$ $al$. ref34 extended the well-known AMP algorithm to propose AMP-Net. These models enjoy the interpretability with speed and tuning-free advantage. But the existing those approaches make little use of the non-local self-similarity pior which plays an important role in image reconstruction.

There have been many previous methods for image reconstruction based on non-local prior. The non-local means (NLM) filter ref35 is highly successful in the image denoising, where it produces a denoised image by calculating the weighted value of the current pixel and its neighbouring pixels. Inspired by NLM, several inverse problem frameworks incorporating non-local regularizer have been proposed ref13 ; ref14 ; ref15 . For instance, Zhang $et$ $al$. ref13 combines the TV regularizer with the non-local regularizer, which is solved using the augmented Lagrangian method and captures non-local features of the image during the iterative process. However, the use of time-consuming NLM filters in the iterations undoubtedly introduces a costly computational complexity. Inspired by deep learning, some recent network-based approaches exploit non-local self-similarity. Liu $et$ $al$. ref36 proposed a network incorporated non-local operations into a recurrent neural network (RNN) for image restoration. Although non-local prior has been widely exploited by both optimization-based and network-based methods, few interpretable deep learning models have introduced this important prior.

This paper combines merit of CS and non-local prior to propose a novel interpretable network, dubbed NL-CS Net. It composed of two parts: the up-sampling module and the recovery module. In the up-sampling phase, we adopted fully connection matrix to stimulate block-wise sampling and initial process. In the recovery phase, we maps the augmented Lagrangian method solving non-local regularized CS reconstruction model into the network, where the network consists of fixed number of phases, each of which corresponds to the one iteration. Rather than the traditional time-consuming NLM operation, the patch-wise non-local network is used to exploit global features. The hyper-parameters involved in NL-CS Net (e.g. sampling matrix, step size, $etc.$) are learned end-to-end, rather than being hand-crafted. Experimented results on natrual images dataset and MRI dataset shows the feasibility and effectiveness of the proposed method compared with the existing methods.

The goal of CS is to reconstruct image from its CS measurement with high quality. Mathematically, given the original image $u\in\mathbb{R}^{N}$, its CS measurements can be obtained by $b=\Phi u\in\mathbb{R}^{M}$, where $\Phi\in\mathbb{R}^{M\times N}$ denotes the sampling matrix and ${M}\diagup{N}$ (${M<<N}$) is commonly regarded as the CS sampling rate. Reconstructing $u$ from $b$ is typically ill-posed. Proposed NL-CS Net combines merit of CS and non-local prior, thus, we first review the traditional optimization-based algorithm to solve the non-local regularized model for CS.

The traditional methods use a preset sampling matrix to recover $u$ from the measured image $b$, which is formulated as solving the following optimization problem:

where $D$ denotes the transform matrix and $R$ is a regularizer, that imposes prior knowledge, such as sparsity and non-local self-similarity.

More effective in suppressing staircase artifact and restoring the detail, optimization-based approaches, combined traditional sparse priors with non-local regularizer, have been proven to achieve superior performance ref13 .

The ${W_{i}}_{j}$ is of the following form:

where $s_{i}$ is the set containing the neighbor of the pixel $i$; ($W_{ij}$) represents the matrix form of NLM; $\alpha$ is a hyper-parameter; $h$ is a controlling factor; the superscript $k$ indicates the number of iterations. In brief, for a given pixel, the NLM filter can be obtined by calculating a weighted average of the surrounding pixels within a search window.

In order to solve Eq. (2), we equivalently transform Eq. (2) into the following problem through variable splitting technique.

where $\omega$ and $x$ are auxiliary variables. Thus, the corresponding augmented Lagrangian function for Eq. (4) is expressed as:

where $\theta$, $\mu$ and $\beta$ are regularization hyper-parameters; $\lambda$, $v$ and $\gamma$ are the Lagrangian multipliers. In this case, the augmented Lagrangian method solves Eq. (5) by the folloing update rule:

By applying the alternating direction method, Eq. (6) can be decomposed into three sub-problems in the following form:

where $d={D^{T}}(\beta D{u^{\rm{(k)}}}-{v^{(k)}}-\beta{\omega^{(k+1)}})-{\gamma^{(k)}}+\theta({u^{(k)}}-{x^{(k)}})+(\mu{\Phi^{T}}(\Phi{u^{(k)}}-\Phi b)-{\lambda^{(k)}})$.

$S(\cdot)$ is a nonlinear shrinkage function with the hyper-parameter $1/{\beta}$, where $S\left({y,z}\right)=sign(y)max(\lvert y\rvert-\frac{1}{z},0)$.

And $\varepsilon$ is the step size of the gradient descent method. The overall algorithm flow is shown in Algorithm 1.

This algorithm uses time-consuming NLM operations in each iteration. It typically requires hundreds of iterations to achieve a satisfactory result, which suffers from a large amount of computation. The transform $D$, sampling matrix and step size are pre-defined, which is very challenging to detain hyper-parameter.

On the basis of Algorithm 1, by combining the merits of optimization-based and network-based approaches, the main idea of this paper is to unfold the solution process of Algorithm 1 into network. The structure of NL-CS Net is shown in Figure 1. For jointly optimizing the sampling matrix and the recovery algorithm, our proposed NL-CS Net consists of a up-sampling module and a recovery module. In the up-sampling phase, we adopted a fully connection matrix to stimulate block-wise sampling and initial process. In the recovery phase, its backbone is designed by mapping the augmented Lagrangian method solving non-local regularized CS reconstruction model into network. The network consists of fixed number of phases, each of which corresponds to the one iteration. Hence, NL-CS Net is composed of $\omega^{(k)}$, $u^{(k)}$ and $x^{(k)}$ modules and update Lagrange multiplier module corresponding to four sub-problem Eqs. (8) , (9) , (10) and (7) sequentially in the $k$-th iteration. We designed a novel model to replace soft-shrinkage function and allowed step size and transform matrix to be learned. The learnable patch-wise non-local method is used to exploit global features, rather than the traditional non-local means operation. The parameters involved in NL-CS Net (e.g. sampling matrix, step size, $etc.$) are learned end-to-end, rather than being hand-crafted.

We validate the proposed model on two tasks: the CS reconstruction of natrual images and MRI images. People are most often in contact with natural images, which is very important. MRI is a non-invasive and widely used imaging technique providing both functional and anatomical information for clinical diagnosis. But long scanning and waiting times may lead to motion artefacts and patient discomfort. MRI acceleration is one of the most successful applications of CS (CS-MRI), which can reconstructs high quality MR images from a few sampling data in k-space. To give quantitative criteria, Peak Signal to Noise Ratio (PSNR) is introduced to analyze the reconstruction performance. We use the Adam optimizer with the default learning rate set to 0.0001 and the batch size to 64. All the network were trained on a workstation configured with Intel Core i7-9700 CPU and RTX2080 GPU, and tested on a workstation configured with Intel Core i7-7820 CPU and GTX1060 GPU.

Inspired by traditional optimization, we proposed a novel CS framework, dubbed NL-CS Net, with the incorporated learnable sampling matrix and non-local piror. The proposed NL-CS Net possesses well-defined interpretability, and make full use of the merits of both optimization-based and network-based CS methods. Extensive experiments show that NL-CS Net have state-of-art performance while maintaining great interpretability. For future work, one direction is to extend our proposed model to other image inverse problems, such as deconvolution and inpainting. Another one is to combine other iterative algorithms with deep learning.

This work was supported by the Natural Science Foundation of Liaoning Province (2022-MS-114).

• •

  Funding: Natural Science Foundation of Liaoning Province (2022-MS-114)

• •

  Conflict of interest/Competing interests: The authors declare that they have no known competing financial interests or personal relationships that could have appeared to influence the work reported in this paper.

• •

  Availability of data and materials: We hereby declare that all data and materials used in this study are publicly available with no restrictions. The data used in this research has been made publicly available and can be accessed directly via https://github.com/bianshuai001/NL-CS-Net.

• •

  Code availability: We declare that the code used in this study is open-source and publicly available for unrestricted use. The code used in this research can be accessed via the link https://github.com/bianshuai001/NL-CS-Net. Anyone can retrieve, download, and use the code for non-commercial purposes, subject to appropriate attribution of the source.