CS-MCNet:A Video Compressive Sensing Reconstruction Network with Interpretable Motion Compensation

In the recent decades, many methods have been proposed to solve the CS reconstruction problem[8],[9],[10],[11]. For image reconstruction of compressive sensing, the algorithm of traditional transformation(e.g., wavelets domain[10],[12])runs fast but with low accuracy. Methods that rely on complex sparsity, such as dictionary learning methods[13]generally have better reconstruction performance but lower computational speed. Furthermore, while most research efforts for image CS problems can be directly applied to video CS tasks, they fail to take advantage of the correlation between adjacent frames in a video sequence.

For video CS recovery, in[8], the authors use Gaussian mixture model(GMM) to recover high-frame-rate videos, and the reconstruction can be efficiently computed as an analytical solution. In[14], the authors propose a motion-compensation and block-based method MC-BCS-SPL, which estimate a motion vector from a reference frame and the under-reconstruction frame can then get prediction to improve recover performance. In general, these traditional methods focus on the design of different priors, transformations and sparsity constraints. However, these methods are usually difficult to determine the hyperparameters, such as thresholds or number of iterations, and due to their computational complexity, they can not perform real-time rebuild.

Driven by the powerful learning capabilities of neural networks, a number of DNN-based approaches have been applied[15],[16],[17],[18],[19]. In[19], the authors cast ISTA into deep network form and develope an effective strategy to solve the proximal mapping associated with the sparsity-inducing regularizer using nonlinear transforms. In[17], the authors propose a fully-connected neural network to reconstruct video temporal CS measurement, and a repetitive pattern measurement mask is proposed to make such a task practical. In[15], the authors propose a network named ”CSVideoNet”. The network combines a multi-rate CNN and a synthesizing RNN to improve the trade-off between compression ratio and spatial-temporal resolution of the reconstructed videos. Compared with iterative algorithms, these feed-forward methods significantly reduce time consumption. However, the structure of these networks is often empirically determined and is ambiguous as a black box, which brings difficulty in making targeted improvements.

To resolve the conflict between interpretability and speed of reconstruction, a technique called algorithm unrolling has been applied recently. The technique was proposed by Gregor et al.[20], and builds neural network by unfolding an iterative optimization algorithm to be a hierarchical architecture, which provides a principled framework by expressing traditional iterative algorithms as neural networks, and offers promise in developing interpretable network. There are several networks using algorithm unrolling to solve CS reconstruction[19],[21], but they are developed for image CS tasks rather than specifically for video CS tasks.

In this work, we develop a network, called CS-MCNet, that attempts to use inter-frame information to improve the reconstruction quality of video CS measurements. By mapping the iterative algorithm MC-BCS-SPL into non-iterative neural network, all of CS-MCNet’s block is designed to correspond to an iteration in MC-BCS-SPL. By end-to-end training, our network can learn all parameters efficiently.

Our main contributions are summarized as follows.1)We use neural network modules to replace the optimization steps in traditional model-based approaches and implement them in a simple form that is easy to realize quickly. 2)We propose a multi-hypothesis motion compensation structure. The module exploits the similarity between neighboring frames to improve the reconstruction quality. To the best of our knowledge, it is the first work that explicitly uses motion compensation for video CS reconstruction in deep neural network. 3)We employ a residual module in the network to further improve performance, and this structure also facilitates the training of deeper neural networks. With these improvements, our work outperforms previous work in terms of both reconstruction quality and computational cost.

The first hidden layer is a fully-connected layer that would provide 3D signal from 2D compressed measurements. Several papers have shown that CNNs can achieve superior performance on CS reconstruction problems compared with simple optimization-based algorithms.[18],[17],[16], thus we use CNN to get the preliminary reconstruction result. Typical CNN architectures used to do recognition, classification, and segmentation are not suitable to the reconstruction problem here. The goal of CNNs in our network is to retain as much detail as possible and need to recover pixels that do not exist based on known information. Therefore, we eliminated the pooling layer, which causes information loss.

To reduce the size of parameters and simplify the network architecture, we use video blocks as input and set the block size to $16\times 16$. The convolutional layers, each of which is followed by a ReLU layer except for the final layer, are carefully designed to get amenable recovery performance. All feature maps are the same size as the reconstructed video block, and the number of feature maps is monotonically reduced. The detailed structure of the CNN is shown in Fig. 3. This process resembles the sparse coding stage in CS, where a subset of dictionary atoms are combined to form an estimation of the original input. To improve final reconstruction performance, we pre-train the CNN before training the whole CS-MCNet, since the path is long from the input to the output of the whole net and pre-training can help prevent the vanishing gradient problem[22].

Traditional video compression algorithms have long exploited motion compensation to improve video-coding quality[23],[24]; In consideration of bit rate at the encoder side, these techniques use single hypothesis in order to limit the amount of motion vector. However, this limitation doesn’t exist since the motion compensation is all calculated at the decoder side of the system. Thus, multi-hypothesis motion compensation can be considered to combine all available best assumptions in the reference frame. A MH CS reconstruction can be represented as an optimal linear combination of all possible reference blocks:

Here, the subscript ”t” and ”i” represents the index of frame in the video and the index of block in the frame, and $H_{t,i}$ is a matrix of dimensionality $B^{2}\times K$ consisted of rasterizations of the possible blocks within the search window in the reference frame, and $K=|H_{t,i}|$. In this context, $w_{t,i}$ represents the linear combination of the columns of $H_{t,i}$; The solution of this optimization can be calculated as a least-squares(LSQ) problem[25].

The proposed network in this paper uses MH motion compensation to improve the recovery performance. Unlike traditional optimization-based solutions to the LSQ problem, we use fully-connected layers to learn the optimal parameters. Due to the similarity of adjacent frames, this MH motion compensation module can be trained appropriately to produce accurate predictions of motion and the recovery quality can be improved by the aggregation of motion and spatial visual features.

To get the reference frame, we design a buffer to store the reconstructed video blocks, as shown in Fig. 2. For the sake of simplicity, we choose to do reconstruction after the reference frame is completely reconstructed. However, the search window actually only involves part of the reference frame, and therefore we can reduce the size of the buffer by carefully designing the rebuilding order. In[15], the authors use LSTM network to do temporal reference, which is similar to motion compensation. However with the experiment, we prove that the utilization of explicit motion compensation module outperforms the RNN based methods and decrease the model size simutaneously.

With the prediction of MH motion compensation, we introduce the residual reconstruction module to further narrow down the gap between $x_{reconstruction}$ and $x$. The output of residual learning is fused with the output of preliminary reconstruction module as the final result.

We get the residual signal $d$ by measuring the result of MH motion compensation and subtracting it from the original measurements $y$,

According to[26], the convolutional layers in residual module could be easily trained to approximate the residual, most of which is zero. As shown in Fig. 4, the residual reconstruction module contains one fully-connected layer and five convolution layers. The fully-connected layer has the same function of recover 3D volume signal from 2D measurements as in preliminary reconstruction module. The rest has decreasing number of feature maps and holds the size of $16x16$. All convolutional layers is followed by a ReLU layer except the last one.

Given the training data pairs $\{(x,x_{ref}\}_{i}$, CS-MCNet firstly gets measurements $y_{i}$ with sensing matrix. Then $y_{i}$ serves as input of the decoder and generates the reconstruction result. We wanted to reduce the discrepancy between the raw measurements and the MH motion compensation measurements, i.e.the residual signal. Therefore, the loss function for CS-MCNet is designed as follows:

where $\lambda$ is the scale factor to control the influence of motion compensation on the total loss. Determined by experiment, we set $\lambda$ to be 0.5 during training to get the best performance.

We choose MSE to calculate the loss, which is a commonly used metric to quantitatively evaluate recovery quality. The proposed framework can also be adapted to other loss functions. Adam optimizer with default parameters is chosen to optimize the proposed network.

As there is no standard dataset designed for video CS, we use UCF-101 dataset[27] to build our own training dataset. UCF-101 dataset includes 13k clips and about 27 hours of video, which is collected from YouTube and is divided into 101 action classes. We extract only the luminance component of the extracted frames and crop the central $160x160$ patch from each frame. All of the patches are segmented into $16x16$ non-overlapping blocks. We randomly choose video sequences from UCF-101 dataset and finally get around 300,000 pairs of data for training and validation in total.

Our model is implemented with PyTorch and all the experiments are performed on a workstation with an Intel Xeon CPU and a Nvidia GeForce RTX2080 GPU. Our networks are trained for 200 epochs with batch size of 400. We normalize the input pre-feature to zero mean and standard deviation one. We set the starting learning rate to 0.01 and divide the learning rate by 10 if the loss of the current epoch is greater than that of the previous epoch. Except for the last subsection, we use 4 cascaded stages in the following experiments.

We compare the reconstruction performance of our proposed method with several reference work of CS reconstruction[14],[11],[8],[15],[19],[17]. The summarized information about all baseline approaches is listed in Table 1. All the methods reconstruct video blocks from its CS measurements independently, and the result of average PSNR, SSIM and time consumption for each method on the test dataset is reported in Fig. 5. To save training time, all methods are tested under CR of 16. From the results we can observe that CS-MCNet outperforms the reference method in terms of metrics and time consumption. Compared with conventional image CS algorithm BCS-SPL, D-AMP, and video CS algorithm MC-BCS-SPL, GMM, our DNN-based CS-MCNet benefits from learnable parameters and feed-forward architecture, and thus gets better reconstruction quality and uses less time. The similar DNN based methods DeepVideoCS, CSVideoCS and ISTANet either uses barely CNN or combines CNN and RNN to extract inexplicit motion features. In contrast to them, our work exploits explicit MH motion compensation, further improving the quality of the reconstruction and compressing the size of the model, which makes it easier and faster to train and deploy.

To further validate the advantages of using MH motion compensation, we do comparison with MC-BCS-SPL and ISTA under different CRs of 4,16 and 64. As shown in Table 2, CS-MCNet achieves relatively better performance, especially under high compression ratio. CS-MCNet outperforms MC-BCS-SPL both on reconstruction quality and time consumption, since its learnable parameters can be better optimized through end-to-end training, and GPU acceleration makes it three orders of magnitude faster than iterative methods. As for ISTANet, which also uses algorithm unrolling and has a similar residual block structure, the utilization of inter-frame information helps our network achieve better reconstruction quality, though the storage of reference frame requires extra storage space and time consumption.

Visual examples of some reference methods are shown in Fig. 6. The ground truth is also shown in Fig. 6. As we can see, CS-MCNet provides the best detail of selected methods and suffers minimal block effect. CS-MCNet produces sharper edges in the highlight areas and a more uniform image overall. This comparison demonstrates that inter-frame information is significant for video CS reconstruction, and the image CS approaches are not suitable for video tasks. Through further optimization, we believe that CS-MCNet has the potential to be applied on real-time reconstruction of high-frame-rate video CS.

In this subsection, we investigate the performance of our proposed network in the presence of measurement noise. The measurement of CS usually involves noise in practice caused by devices, and the measurement model should now be modified as

where n is the additive measurement noise.

We conduct experiment with input measurements contaminated by random Gaussian noise, and all other parameters remain the same as in the noiseless case. We test the performance at four level of SNR from 20dB to 50dB under CR of 16. The result is shown in Fig. 7. It can be observed that at different noise level, CS-MCNet can achieve stable reconstruction performance and outperform the reference methods consistently. It is worth emphasizing that we did not retrain the network with noise measurement and all the experiments in this subsection are implemented with noiseless model, which shows the robustness of our model under different measurement conditions. Besides, it is easy to combat with performance degradation under noise by cascading our network with a deep denoising architecture or other denoising algorithm.

As described earlier, the number of stages of CS-MCNet corresponds to the number of iterations of the original algorithm. In this subsection, we mainly focuses on the structure parameter of CS-MCNet, i.e. the nubmer of stages. From Table 3, we can find that as the number of stages increases from 2 to 4, the performance improves under different CRs. This can be explained by the fact that the deeper the neural network, the better its learning capacity. However, as the number of stages increases to 5, the performance deteriorates. We speculate that the reason is that while a deeper structure may help to fit the training data more accurately, it also makes it more difficult to train, resulting in an undertrained model. The training time increased slightly when the number of stages is less than 4 and increases rapidly when the number of stages is further increased. For CS-MCNet with no more than 4 stages, the time used to train an epoch varies from 30 to 40 minutes with GPU acceleration, but for networks with 5 or more stages, it can take more than an hour. To strike a balance between the effectiveness of the reconstruction and the cost of training, we empirically decides to use CS-MCNet with 4 stages.