FSOINet: Feature-Space Optimization-Inspired Network for Image Compressive Sensing

In the BCS [2] strategy, an image $\bm{X}\in\mathbb{R}^{H\times W}$ is divided into $\frac{H}{\sqrt{N}}\times\frac{W}{\sqrt{N}}$ non-overlapping image blocks with size of $\sqrt{N}\times\sqrt{N}$, then each of them is reshaped into a vector $\bm{x}\in\mathbb{R}^{N}$. When the CS sampling rate is $r$, the same sampling matrix $\Phi\in\mathbb{R}^{M\times N}$ is used to sample each $\bm{x}$, where $M=\left\lfloor r\times N\right\rfloor$. In addition, Cui et al. [14] proposes to use a learnable sampling matrix instead of fixed random Gaussian Matrix to get CS measurements, and have achieved gratifying results. In this paper, we use a sampling sub-network which is denoted as $\mathcal{F}_{\Phi}(\cdot)$ to sample the image and use convolution without bias to simulate the block sampling process described above. To be specific, we set $\Phi$ as learnable network parameters, reshape each row of $\Phi$ into a convolution kernel of size $1\times\sqrt{N}\times\sqrt{N}$ to obtain $W_{\Phi}$. Therefore, the sampling process can be described by convolution with a stride of $\sqrt{N}$ as:

To obtain a reasonable initial estimate of each image block from CS measurements and not introduce more learnable parameters in the meanwhile, we use $\Phi^{\mathrm{T}}$ to complete the initial reconstruction. Like the sampling process, each row of $\Phi^{\mathrm{T}}$ is reshaped into a convolution kernel of size $M\times 1\times 1$ to obtain $W_{\Phi}^{\mathrm{T}}$, then PixelShuffle is used to obtain the initial reconstruction image, which can be illustrated as:

In previous optimization-inspired ICS networks, each phase implements supplementing information in pixel space according to Eq.2 and sends the denoised image to the next phase, although the denoising is finished in feature space. This process makes the information flow in pixel space phase by phase, which limits the use of the robust feature representation ability of CNN. In this paper, we propose the idea of implementing the information flow phase by phase in feature space, including supplementing the information from measurements and dual-scale denoising all in the feature space, to fully utilize the feature representation ability of CNN and retain more details when the information flows to the next phase. The deep reconstruction subnet is shown in Fig.1.

Information supplement in feature space: Consistent with the optimization algorithm, we also use measurements throughout the entire reconstruction process to guarantee the solution complies with the degradation process $\bm{y}=\Phi\bm{x}$. And different from previous optimization-inspired ICS methods, we propose the idea of supplementing information in feature space instead of that in pixel space from measurements. To implement this idea, we build a Feature-space Information Supplementing Module (FSIM) to map the gradient of the fidelity term $\frac{1}{2}\|{{\Phi}\bm{x}-\bm{y}}\|^{2}_{2}$ in Eq.1 to the feature space and fuse it with the image features, which is shown in Fig.2(a). It is worth noting that $\mathcal{F}_{\Phi}(\cdot)$ and $\mathcal{F}_{\Phi}^{\mathrm{T}}(\cdot)$ in Fig.2(a) are block operations, which will introduce blocking artifacts. Here, ResBlock operates on the entire image, which can suppress blocking artifacts.

Dual-scale denoising in feature space: Benefiting from the data-driven characteristics of deep learning, the inherent prior knowledge of the image can be learned from the training set to denoise the image. In this paper, we propose a dual-scale structure network to improve denoising efficiency. And we design a neural network module called Dual-scale Denoising Module (DDM) to implement it, as shown in Fig.2(b). To maintain not too many parameters, only two ResBlocks are used to denoise the high-resolution features and low-resolution features respectively in DDM. Then, the denoised features are fused to obtain the output features, which will be sent into the next phase. Finally, we implement the information flow in feature space phase by phase. The high-resolution features maintain the same resolution as the original image, while the low-resolution features are subsampled from the high-resolution features with stride-2 Conv and have twice the number of channels.

For an original image $\bm{X}$, the proposed model first obtains CS measurements $\bm{Y}$ by sampling it and then predicts the reconstructed image $\bm{X}_{\text{rec}}$. We optimize our FSOINet end-to-end through the following loss function:

where $\mathcal{L}_{mse}$ is the Mean Squared Error (MSE) of the original image $\bm{X}$ and the reconstructed image ${X}_{\text{rec}}$. In addition, $\mathcal{L}_{orth}$ is the orthogonal constraint of the sampling matrix in [12]:$\mathcal{L}_{orth}(\Phi)=\frac{1}{M^{2}}\|{{\Phi\Phi^{\mathrm{T}}}-\text{I}}\|^{2}_{F}$, where I represents the identity matrix. The regularization parameter $\gamma$ in Eq.6 is set to 0.01 in our experiments.

For training, following [5], we use 400 images from the training set and test set of the BSDS500 dataset [15]. The training images are cropped to 89600 $96\times 96$ pixel sub-images with data augmentation. For the network parameters, the block size $\sqrt{N}$ is 32, the channel number $C$ is 16, the phase number $N_{k}$ is 16 and the batch size is 32. The size of the unspecified convolution kernel is $3\times 3$. We use Adam [16] optimizer to train the network with the initial learning rate of $2\times 10^{-4}$, which is decreased to $5\times 10^{-5}$ through 100 epochs using the cosine annealing strategy [17] and the warm-up epochs are 3. All the experiments are implemented in PyTorch with a Core i5-6500 CPU Intel and a GTX2080ti GPU. For testing, we utilize three widely-used benchmark datasets, including Set11 [3], BSDS68 [18] and Urban100 [19]. Color images are processed in the YCbCr space and evaluated on the Y channel. Two common-used image assessment criteria, Peak Signal to Noise Ratio (PSNR) and Structural Similarity (SSIM), are adopted to evaluate the reconstruction results.

Considering that the sampling matrix obtained by training is often better than the fixed Gaussian random matrix. In order to make fair comparisons, we compare our FSOINet with five state-of-the-art ICS methods whose sampling matrixes are learnable too, including CSNet [5], SCSNet [6], SPLNet [10], AMPNet [11], and OPINE-Net [12]. Extensive experiments have proved the advantages of FSOINet in terms of quality and visualized results.

Table 1 and Fig.3 clearly show that our FSOINet outperforms all the other competing methods at all the CS sampling rates. It is worth noting that the three optimization-inspired methods, SPLNet, OPINENet, and AMPNet, are significantly better than vanilla neural networks CSNet and SCSNet in terms of reconstruction quality. But our FSOINet got a more compelling reconstruction quality and reconstruct sharper edges and clearer background at all the sampling rates.

Table 2 provides a comparison of model sizes and calculation times for different methods at a CS ratio of 0.5. Our model has the fewest parameters and medium calculation speed while achieving the best reconstruction results.

This section mainly analyzes the effectiveness of using FSIM to replace the gradient descent in the optimization algorithm. We retrain our network without FSIM, then the network is similar to vanilla neural networks, which only use the measurements during the initial reconstruction, dubbed VNet. At the same time, we change the FSIM to the gradient descent in the pixel domain and retrain our network. At this time, the network which we named OINet is more similar to the previous optimization-inspired neural networks. For a fair comparison, we move the convolution block used in FSIM to DDM to maintain similar model parameters.

As shown in Table 3, the PSNR/SSIM of OINet is 0.54dB/0.0076 lower than that of FSOINet at the sampling ratio of 0.1, which reflects the effectiveness of mapping the gradient information to the feature space. While VNet’s PSNR/SSIM is 0.57db/0.0047 lower than that of OINet, which indicates the insights of the optimization algorithm could help the neural network structure design.