Practical Compact Deep Compressed Sensing

TL;DR

Extensive experiments on natural image CS, quantized CS, and self-supervised CS demonstrate the superior reconstruction accuracy and generalization ability of PCNet compared to existing state-of-the-art methods, particularly for high-resolution images.

Abstract

Recent years have witnessed the success of deep networks in compressed sensing (CS), which allows for a significant reduction in sampling cost and has gained growing attention since its inception. In this paper, we propose a new practical and compact network dubbed PCNet for general image CS. Specifically, in PCNet, a novel collaborative sampling operator is designed, which consists of a deep conditional filtering step and a dual-branch fast sampling step. The former learns an implicit representation of a linear transformation matrix into a few convolutions and first performs adaptive local filtering on the input image, while the latter then uses a discrete cosine transform and a scrambled block-diagonal Gaussian matrix to generate under-sampled measurements. Our PCNet is equipped with an enhanced proximal gradient descent algorithm-unrolled network for reconstruction. It offers flexibility, interpretability, and strong recovery performance for arbitrary sampling rates once trained. Additionally, we provide a deployment-oriented extraction scheme for single-pixel CS imaging systems, which allows for the convenient conversion of any linear sampling operator to its matrix form to be loaded onto hardware like digital micro-mirror devices. Extensive experiments on natural image CS, quantized CS, and self-supervised CS demonstrate the superior reconstruction accuracy and generalization ability of PCNet compared to existing state-of-the-art methods, particularly for high-resolution images. Code is available at https://github.com/Guaishou74851/PCNet.

Figures (12)

• Figure 1: Illustrations of the proposed PCNet architecture and a typical single-pixel hardware implementation duarte2008single for CS acquisition. (a) Following zhang2020optimizationchen2022content, our PCNet consists of a sampling subnet (SS), an initialization subnet (IS), and a recovery subnet (RS). The SS simulates the sampling process, the IS performs the transformation from measurement to image, while the RS, which is algorithm-unrolled and transmission-augmented, further refines the initialization from the IS to provide a final reconstruction. (b) In practice, linear sampling matrices (operators) are converted into a global matrix form as shown in Fig. \ref{['fig:sampling_operator_comparison']} (a), of size $M\times N$. This matrix is then split and re-organized into $M$ separate modulating patterns (each corresponding to a matrix row of size $H\times W$) to be individually loaded onto specific spatial light modulators (SLMs) like the digital micro-mirror devices (DMDs), with a spatial resolution of $N=H\times W$, for sequential sampling.
• Figure 2: Illustrations of different CS sampling operators. (a) Global dense matrix linearly aggregates all the $N$ pixel values of image $M$ times. (b) Block-diagonal matrix divides image into non-overlapping blocks and performs sampling block-by-block. (c) Deep sampling NN zheng2020sequentialfan2022global uses stacked convolutions to gradually reduce the data volume and obtain compressed measurements. (d) Our collaborative sampling operator (COSO) consists of a deep filtering step and a fast global sampling step. The former (denoted by $\mathcal{T}$) contains seven convolution layers and is modulated by the CS ratio. It learns an implicit compact representation of matrix into lightweight $\mathcal{T}$ from data. The latter is achieved through a dual-branch transformation and masking. More details are illustrated in Fig. \ref{['fig:sampling_operator']} (a)-(c). Note that in CS imaging systems like the one shown in Fig. \ref{['fig:arch']} (b), any sampling operator must be converted to its global matrix form as in (a) for hardware deployment. (e) We further provide a practical scheme (see Algo. \ref{['alg:matrix_extraction']}) to extract the explicit matrix from any given (black-box) linear system for sampling matrix deployment and analysis.
• Figure 3: Our pilot experiments evaluating five sampling matrices on T91 dong2015image (left) and Train400 zhang2017beyond (right). These matrices include global dense (G. D.) Gaussian, DCT, Hadamard, block-diagonal (B.) Gaussian, and scrambled block-diagonal (S. B.) Gaussian matrices. All images are center-cropped to 256$\times$256, and the block size $B$ is set to 32. (1) Top: For an image set $\mathcal{S}=\left\{\mathbf{x}_i\right\}$, the restricted isometry constant of sampling matrix $\mathbf{A}$ is estimated by $\delta \approx \max_i \left\{abs({\lVert \mathbf{Ax}_i \rVert_2^2}/{\lVert \mathbf{x}_i \rVert_2^2}-1)\right\}$. Small $\delta$ values are preferred as they indicate that $\mathbf{A}$ nearly preserves orthogonality within the domain of $\mathcal{S}$. (2) Bottom: The average absolute measurement value of $\lvert \mathbf{y}_i \rvert=\lvert \mathbf{A}_i\mathbf{x}\rvert$ in logarithmic scale is calculated by $\frac{1}{j}\Sigma_j\log \lvert \mathbf{A}_i\mathbf{x}_j \rvert$ at the point of $i=\gamma N$ for all possible CS sampling rates $\gamma\in[0,1]$. A large $\lvert \mathbf{y}_i \rvert$ value suggests a high correlation between the $i$-th matrix row $\mathbf{A}_i$ and the images in $\mathcal{S}$.
• Figure 4: Illustration of two block-based CS sampling schemes for an image of size $N=H\times W$, block of size $n=B^2$, and rate $\gamma=m/n$. For simplicity, we only present the simulated sampling process for one block, setting $H=W=6$ and $B=m=2$ in our toy example. (a) The block-diagonal matrix gan2007block divides the image into non-overlapping $B\times B$ blocks and obtains their measurements one-by-one. This operation can be efficiently simulated by a convolution layer (left). (b) The introduced scrambled block-diagonal matrix do2008fast randomly shuffles all pixels before its block-by-block sampling process (middle). This approach is equivalent to non-overlapping sampling performed on the whole image (right), thus benefiting from a global receptive field and low complexity.
• Figure 5: Illustration of our proposed collaborative sampling operator (COSO) $\mathcal{G}_\mathbf{A}$ for sampling an image with size $N=H\times W$ and rate $\gamma$. (a) The input $\mathbf{X}$ is first adaptively filtered by a CS ratio-aware lightweight CNN $\mathcal{T}$, which consists of seven convolution layers and is conditioned by the rate vector $\mathbf{z}=\left[\gamma_D,\gamma_G\right]$ satisfying $\gamma=\gamma_D+\gamma_G$, to generate two smoothed results $\mathbf{X}_D$ and $\mathbf{X}_G$. When the whole sampling operator is utilized in the reconstruction network, we introduce two supplementary paths into $\mathcal{T}$ (indicated by the red dotted lines with arrows) to enable a high-throughput feature-level information flow. (b) The DCT-based fast dense sampling branch (D-branch) transforms $\mathbf{X}_D$ into the DCT domain and discards the last $(1-\gamma_D)N$ coefficients in Zig-Zag order to obtain the measurement $\mathbf{Y}_D$. (c) The scrambled block-diagonal Gaussian sampling branch (G-branch) randomly permutes all the pixels of $\mathbf{X}_G$ and performs a channel-wise masked convolution to simulate block-diagonal sampling and produce measurement $\mathbf{Y}_G$ at the ratio of $\gamma_G$. (d) For practical deployment and analysis, all components of our operator, composed of (a)-(c), can be merged into a simple equivalent linear pipeline, $\mathbf{y=M(\Phi x+b^\prime)=(M\Phi) x+(Mb^\prime)}$, by applying the matrix distributive and associative laws.