RDFNet: Regional Dynamic FISTA-Net for Spectral
Snapshot Compressive Imaging

Our work is related to the recent dynamic mechanism. In particular, Chen et al. [31] propose a dynamic convolution that aggregates multiple convolution kernels dynamically based on the input. Brabandere et al. [32] present a dynamic filter network to dynamically generate position-specific filters on pixel inputs. CondConv [33] generates convolution kernels by combining several filters through a routing function. Recently, PAConv [34] develop a position adaptive convolution operator with dynamic kernel assembling for point cloud processing. However, the region-based dynamic mechanism has not yet been explored in the field of SCI reconstruction. Zhang et al. [35] design a weight for each pixel in an image, use the same transformation to perform super-resolution, and add the weights to obtain a mixed transformation for the entire image. In comparison, we split the image into regions instead of pixels and perform distinct domain transformations with pixel-level adaptive thresholding for different regions while retaining neighborhood information.

With the rise of neural networks, some algorithms try to use the convolution to obtain the DIP but there is no deep network structure, forming a machine learning algorithm.

Bacca J. et al [36] proposed a network for spectral reconstruction without training according to the ideas of solving ill-posed problems with low rank. It is mainly achieved by analyzing the low-rankness of images at the first layer of the network. Evaluating the difference between minimized compression measurements and predictions by the use of $l_{2}$. However, it does not really use a deep neural network in the process of solving, but uses several convolutions to help getting the prior of recon. Van Veen D. et al [37] also proposed an untrained model, which may belong to a kind of machine learning method. The neural network is only used to learn the weight of the prior information, not the way to really obtain the prior information, and the neural network here is not deep, but only uses the volume product. Inspired by the linear mixture model (LMM) for spectral image, Gelvez T. et al [38] decomposed the image into a matrix, and uses the neural network to learn the weights and features of each matrix as the depth prior of the image for reconstruction.

In [39], DIP is employed as a refinement process of the trained network for the reconstruction of a single image. The other related work is DeepRED [40], where DIP is combined with Regularization by Denoising (RED) [41]. And the hyperspectral way of using deepred [40] is [42]. In fact, most of the processes have nothing to do with the design of deep neural networks, and does not using the characteristics of adjusting the transformation domain for optimizing reconstruction tasks.

Inspired by the prevalence of deep learning in the field of high-level visions, some researchers have attempted to use deep convolution neural networks (CNNs) to learn the inverse process. These deep learning-based algorithms can be divided into three streams: End-to-End (E2E) [43, 24, 44, 22], Plug-and-Play (PNP) [45, 46, 21], and deep unfolding [47, 26, 48, 23, 49].

Most recently, the GAP-net proposed in [47] has achieved good results in both video and spectral SCI. A deep unfolding based on the Gaussian scale mixture model has been developed in [48] for spectral SCI reconstruction. DNU [23] has contributed to the introduction of a new prior for optimization. Zhang and Wang [49] first learned the tensor low-rank prior of hyperspectral images in the feature domain by DNN to promote the reconstruction quality.

Nonetheless, these methods still show limitations in modeling sparsity representations. Besides, the guidance of regional characteristics for adjusting the reconstruction transformation domain is under-studied.

Spectral snapshot compressive imaging (SCI) sysmtem comprises of a hardware encoder and a software decoder. The encoder denotes the optical system that compresses 3D data cube $(x,y,\lambda)$ to a snapshot measurement on a 2D detector. The decoder denotes the reconstruction algorithm used to recover the 3D data cube from the snapshot measurement. Here, we focus on the coded aperture snapshot spectral imaging (CASSI) system that uses a fixed mask and a disperser to implement band-wise modulation.

As shown in Fig.2, each spatial position of the scene is modulated by a coded aperture (mask) that blocks or unblocks the incoming light. Then the coded spectral scene passes through the prism to introduce a horizontal shifting. Finally, the coded shifted spectral scene is integrated along the spectral axis by the detector, resulting in 2D compressed measurement.

Following the theory in [52], let $X\in\mathbb{R}^{N_{x}\times N_{y}\times N_{\lambda}}$ denote the 3D spatial-spectral cube and $M^{0}\in(0,1)^{N_{x}\times N_{y}}$ the physical mask used for signal modulation. We use $X^{\prime}\in\mathbb{R}^{N_{x}\times N_{y}\times N_{\lambda}}$ to represent the modulated signal where images at different wavelengths are modulated separately by the same mask. For $n_{\lambda}\in\left\{1,....,N_{\lambda}\right\}$, we have:

where $\odot$ represents element-wise multiplication.

Next comes the disperser, which disperses the light to different spatial locations based on their wavelengths. After the modulated cube passes the disperser, $X^{\prime}$ is tilted and considered to be sheared along the $y-$axis. We use $X^{\prime\prime}\in\mathbb{R}^{N_{x}\times(N_{y}+N_{\lambda}-1)\times N_{\lambda}}$ to denote the tilted cube and assume $\lambda_{c}$ to be the reference wavelength. That is, image $X^{\prime}(:,:,n_{\lambda_{c}})$ is not sheared along the $y-$axis, we hence have:

where $(u,v)$ indicates the coordinate system on the detector plane, and $\lambda_{n}$ is the wavelength of channel $n_{\lambda}$. Here, $d(\lambda_{n}-\lambda_{c})$ signifies the spatial shifting for channel $n_{\lambda}$. The compressed measurement at the detector $y(u,v)$ can thus be modeled as

since the sensor integrates all the light in the wavelength range $[\lambda_{min},\lambda_{max}]$, where $f^{\prime\prime}$ is the continuous representation of $F^{\prime\prime}$. In discretized form, the captured 2D measurement $Y\in\mathbb{R}^{N_{x}\times(N_{y}+N_{\lambda}-1)}$ is

which is a compressed frame containing information of all the modulated spectral channels and $\epsilon\in\mathbb{R}^{N_{x}\times(N_{y}+N_{\lambda}-1)}$ represents the measurement noise. For simplicity purpose, we denote $M\in R^{N_{x}\times(N_{y}+N_{\lambda}-1)\times N_{\lambda}}$ as the shifted version of the physical mask corresponding to different wavelengths,

Similarly, for each signal frame at different wavelengths, the shifted version $\tilde{X}\in R^{N_{x}\times(N_{y}+N_{\lambda}-1)\times N_{\lambda}}$ is

Based on the above, measurement $Y$ can be represented as

This corresponds to the encoding process of SCI in Fig.2. Note that the 3D mask $M$ can be obtained by calibration. Given the solved $\tilde{X}$, we can obtain the desired 3D cube by shifting it back to $F$ based on the relationship in Eq.(6),

where $x_{k}=vec(X_{k})$ represents the vectorization of frame $k$ and $D_{k}=Diag(vec(M_{k}))$ is a diagonal matrix with diagonal elements vectorized of $M_{k}$. We obtain the forward model

which is the core problem of spectral SCI reconstruction. Conventional methods [53, 54, 55] usually employ a regularization term $R(\bm{x})$ as prior to constrain the solution in desired signal space. These algorithms aim to find an estimated $\bm{\bar{x}}$ of $\bm{x}$ by solving the following problem:

where $\lambda$ is a parameter to balance between the fidelity and the regularization term. Eq.(10) is usually solved by iterative algorithms with various image priors of $R(\bm{x})$ including sparsity [53], total variation [54], deep denoising prior [45, 21], autoencoder prior[43], etc.

Given the measurement $\bm{y}$ and the modulate mask $\bm{\Phi}$, the problem of reconstructing hyperspectral image $\bm{x}$ can be solved by LASSO optimization [56]. Using $l_{1}$-norm to impose sparsity constraint for coefficients [25], the reconstruction problem in Eq.(10) can be converted as

where $\bm{\Psi x}$ denotes the coefficients in the transformation domain. By introducing an auxiliary parameter $\bm{r}^{k}$, the unconstrained optimization in Eq.(11) can be solved by iterative steps [27]:

where $k\geq 1$, $\bm{z}^{1}=\bm{x}^{0},t^{1}=1$, $\rho$ represents the step size. $\bm{z}^{k+1}$ is a new strating point for next iteration. In each step, we directly utilize the updated $t^{k}$ to calculate $\bm{z}^{k}$.

Conventional FISTA algorithm [27] regards the image as a whole and performs simple global static transformation using a single mapping function. However, we found that there exist significant differences among different regions in a measurement. Hence we are dedicated to applying distinct mapping functions for different regions based on their unique characteristics to realize region-adaptive transformation.

Based on the tensorlization operations [25], we re-reference the theoretical process of RDFNet to match our regional dynamic transformation and pixel-wise soft-thresholding. Specifically, we divide input $\bm{r}^{k}$ into a series of regions $\bm{r}_{i}^{k}\in\left\{\bm{r}_{1}^{k},...,\bm{r}_{M}^{k}\right\}$, and process each individual region using a dynamic mapping function determined by the region’s characteristics. By using $F(\cdot)$ to learn the sparsest representation of spectral images, we can obtain the regional results $\bm{x}_{i}^{k}\in\left\{\bm{x}_{1}^{k},...,\bm{x}_{M}^{k}\right\}$ based on the relationship in Eq.(11):

Following the Parseval Theorem

where $\bm{D}$ is an orthonormal transformation matrix. Eq.(16) can be converted as

According to the soft-thresholding theory [57], we adopt soft-thresholding operator to obtain the closed-form solution for each region:

where $\hat{F}^{\prime},\hat{F}$ are the mixture dynamic transformation. By summing up the regional results $\bm{x}_{i}^{k}\in\left\{\bm{x}_{1}^{k},...,\bm{x}_{M}^{k}\right\}$, we can obtain the final solution $\bm{x}^{k}$.

To achieve the adjustment of regional dynamical transformation $(\hat{F}(\cdot),\hat{F}^{\prime}(\cdot))$, we first design multiple mapping functions $\left\{(F_{i}(\cdot),F^{\prime}_{i}(\cdot))\right\}_{i=1}^{N}$ to represent different fundamental transformations. Then we derive several regional characteristic-driven weights $\bm{w}_{i}^{k}$ corresponding to each mapping function. Hence the transformation can be dynamically adjusted according to the weights. The solution can be calculated as:

For soft-thresholding, we aim to learn an adaptive threshold for each pixel within a region. Specifically, we design a pixel-wise adaptive soft-thresholding $soft^{\bm{\tau}_{i}^{k}}$ by using $sgn(\bm{x})$ to shrinkage every signal pointed among transformations:

where $\bm{\tau}_{i}^{k}$ is the adaptive soft thresholding determined by regions’ characteristic $\bm{r}_{i}^{k}$.

Inspired by the skip connection in ResNet [58], we obtain the closed-form solution of Eq.(18) as:

Hence, we can achieve the desired solution in a learnable manner.

Next, we design a novel Regional Dynamic Network (RDFNet) to implement the above regional dynamic FISTA algorithm.

Specifically, we first slide a $H\times W$ extraction window on the input 2D measurement of size $H\times(W+L-1)$ with slide step of one pixel, and split the input into $L$-channel image of size $H\times W$. Then the split sub-images are fed into the tensorizing pretreatment as stated in [25] to transfer the iteration from vector to tensor form to reduce interference time and memory footprint.

Based on the deep-unfolding framework, we propose a novel deep architecture for solving the proximal mapping problem of compressive sensing reconstruction by using a dynamic nonlinear sparsifying transformation at each iterative phase. It contains three main components.

The proposed regional weighting module extracts regional characteristics to generate a region-wise dynamic weight to guide the optimal sparsifying transformation for each region. Different from previous iterative methods that perform a single fixed transformation for the entire image, the developed multiple dynamic blocks aim at learning different fundamental transformations for different regions by exploiting their corresponding unique characteristics. Besides, an adaptive threshold module is designed in each dynamic block to learn pixel-wise adaptive soft-thresholds.

Finally, we merge the outputs of multiple dynamic blocks via a summation of region-based weights to obtain the final output.

Obviously, the number of dynamic blocks $N$ plays an important role. A larger $N$ contributes to more diversified domain transformation for sparse representation. Nevertheless, too many transformation domains may lead to redundancies and cause heavy memory and computational overhead. We find that setting $N=3$ is appropriate, which is discussed in Sec.  IV-D b)

Given the training data pairs ${(\bm{y}_{j},(\bm{x}_{gt})_{j})^{D}_{j=1}}$, RDFNet takes the measurement $\bm{y}$ as input and generates the reconstruction $\bm{x}$. We seek to reduce the discrepancy between $\bm{x}$ and $(\bm{x}_{gt})$, which indicates the accuracy of the inverse function, while satisfying the symmetry constraint in each dynamic block. Furthermore, we measure the sparsity of spectral frames in the learned domain. For the output $\bm{X}^{k}$ in the $k$-th phase, denote $\bm{X}_{gt}$ as the tensor form of the groundtruth $\bm{x}_{gt}$, we design the loss function for RDFNet as:

The final loss is a weighted combination of the above three terms:

where $\alpha$, $\beta$ and $\gamma$ are balancing coefficients. By default, we set $\alpha=1$, $\beta=0.01$ and $\gamma=0.001$.

We unfold the proposed iterative algorithm into five phases. Each phase contains one RDFNet. All experiments are conducted on a NVIDIA RTX-3090. We set the number of dynamic block as 3 and the regional pooling kernel size as $5\times 5$. We train the model for $3,000$ epochs using Adam optimizer [19] with learning rate 0.0001 and batch size 4. The Peak-Signal-to-Noise Ratio (PSNR) and structural similarity index (SSIM) [61] are employed to evaluate the quality of reconstructed spectral data-cube.

We test our methods on real SD-CASSI data [64, 22] that captures real scenes with $28$ wavelengths ranging from $450\rm{nm}$ to $650\rm{nm}$ and has 54-pixel dispersion in the column dimension. Thus, the measurements captured by the system have a spatial size of $660\times 714$. Fig. 6 shows the reconstruction results of scene1 with four channels by RDFNet and other competing methods. One can observe that our method well recovers textures in both spectral and spatial dimensions.

As shown in Fig. 5 (a), the quality of reconstruction (evaluated by PSNR and SSIM) is gradually increasing. Our RDFNet achieves the best performance and outperforms the FISTA-Net baseline by $2.79\rm{dB}$ in average PSNR. Tab. V shows the improvements by separately incorporating the dynamic block and adaptive soft-thresholding are $0.3\rm{dB}$ and $1.7\rm{dB}$ in average, respectively. It indicates that the regional dynamic strategy largely contributes to the performance gain and the adaptive soft-thresholding brings additional improvement.