Coarse-to-Fine Sparse Transformer for Hyperspectral Image Reconstruction

Introduction

Hyperspectral images (HSIs), which contain multiple continuous and narrow spectral bands, can provide more detailed information of the captured scene than normal RGB images. Based on the inherently rich and detailed spectral signatures, HSIs have been widely applied to many computer vision tasks and graphical applications, $e.g.,$ image classification , object tracking , remote sensing , medical imaging , $etc$.

To collect HSI cubes, traditional imaging systems scan the scenes with multiple exposures using 1D or 2D sensors. This imaging process is time-consuming and limited to static objects . Thus, conventional imaging systems cannot capture dynamic scenes. Recently, researchers have developed several snapshot compressive imaging (SCI) systems to capture HSIs, where the 3D HSI cube is compressed into a single 2D measurement . Among these SCI systems, coded aperture snapshot spectral imaging (CASSI) stands out as a promising solution and has become an active research direction . CASSI systems modulate HSI signals at different wavelengths by a coded aperture (physical mask) and then vary the modulation by a disperser, i.e., to shift the modulated images at different wavelengths to different spatial locations on the detector plane. Subsequently, a reconstruction algorithm is used to restore the 3D HSI cube from the 2D compressive image, which is a core task in CASSI.

To solve this ill-posed inverse problem, traditional methods mainly depend on hand-crafted priors and assumptions. The main drawbacks of these model-based methods are that they need to tweak parameters manually, leading to poor generality and slow reconstruction speed. In recent years, deep learning methods have shown the potential to speed up the reconstruction and improve restoration quality for natural images . Hence, convolution neural networks (CNNs) have been used to learn the underlying mapping function from the 2D compressive measurement to the 3D HSI signal. Nevertheless, these CNN-based methods yield impressive results but show limitations in capturing non-local self-similarity and long-range dependencies.

In the past few years, the natural language processing (NLP) model Transformer has gained much popularity and achieved great success in computer vision. Transformer provides a powerful model that excels at exploring global inter-dependence between different regions to alleviate the constraints of CNN-based methods. Nonetheless, directly applying vision Transformers to HSI reconstruction encounters two main issues that cannot be ignored. Firstly, the HSI signals exhibit high spatial sparsity as shown in Fig. 1 (a). Some dark regions are almost uninformative. However, previous local or global Transformers process all spatial pixel vectors inside non-overlapping windows or global images into tokens without screening and then feed the tokens into the multi-head self-attention (MSA) mechanism. Many regions with limited information are sampled, which dramatically degrades the model efficiency and limits the reconstruction performance. Secondly, previous Transformers linearly project all the tokens into $query$, $key$, and $value$, and then perform matrix multiplication for calculating MSA without clustering. Yet, some of the tokens are not related in content. Attending to all these tokens at once lowers down the cost-effectiveness of model and may easily lead to over-smooth results . Besides, the computational complexity of global Transformer is quadratic to the spatial dimensions, which is nontrivial and sometimes unaffordable. MST calculates MSA along the spectral dimension, thus circumventing the HSI spatial sparsity.

Hence, how to combine HSI sparsity with learning-based algorithms still remains under-explored. This work aims to investigate this problem and cope with the limitations of existing CNN-based and Transformer-based methods.

In this paper, we propose a novel method, coarse-to-fine sparse Transformer (CST), for HSI reconstruction. Our CST composes two key techniques. Firstly, due to the large variation in HSI informativeness of spatial regions, we propose a spectra-aware screening mechanism (SASM) for coarse patch selecting. To be specific, in Fig. 1 (a), our SASM partitions the image into non-overlapping patches and then detects the patches that are informative of HSI representations. Subsequently, only the detected patches (yellow) are fed into the self-attention mechanisms to decrease the inefficient calculation of uninformative regions (green) and promote the model cost-effectiveness. Secondly, instead of using all projected tokens at once like previous Transformers, we aim to calculate self-attention of tokens that are closely related in content. Toward this end, we customize spectra-aggregation hashing multi-head self-attention (SAH-MSA) for fine pixel clustering as shown in Fig. 1 (b). SAH-MSA learns to cluster tokens into different groups (termed $buckets$ in this paper) by searching similar elements that produce the max inner product. Tokens inside each $bucket$ are considered closely related in content. Then the MSA operation is applied within each $bucket$. Finally, with the proposed techniques, we enable a coarse-to-fine learning scheme that embeds the HSI spatial sparsity into learning-based methods. We establish a series of small-to-large CST families that outperform state-of-the-art (SOTA) methods while requiring much cheaper computational costs.

The main contributions of this work can be summarized as follows:

We propose a novel Transformer-based method, CST, for HSI reconstruction. To the best of our knowledge, it is the first attempt to embed the HSI spatial sparsity nature into learning-based algorithms for this task.

We present SASM to locate informative regions with HSI signals.

We customize SAH-MSA to capture interactions of closely related patterns.

Our CST with much lower computational complexity significantly surpasses SOTA algorithms on all scenes in simulation. Moreover, our CST yields more visually pleasant results than existing methods in real HSI restoration.

Related Work

Conventional HSI reconstruction methods rely on hand-crafted image priors. For instance, gradient projection algorithms are exploited to handle the HSI sparseness. In addition, total variation regularizers are employed by GAP-TV while the low-rank property and non-local self-similarity are used in DeSCI . Nonetheless, these traditional model-based methods suffer from low reconstruction speed and poor generalization ability. Recently, CNNs have been used to solve the inverse problem of spectral SCI. These CNN-based algorithms can be divided into three categories, i.e., end-to-end (E2E) methods, deep unfolding methods, and plug-and-play (PnP) methods. E2E algorithms apply a deep CNN as a powerful model to learn the E2E mapping function of HSI restoration. Deep unfolding methods employ multi-stage CNNs trained to map the measurements into the desired signal. Each stage contains two parts, i.e., linear projection and passing the signal through a CNN functioning as a denoiser. PnP methods plug pre-trained CNN denoisers into model-based methods to solve the HSI reconstruction problem. Nonetheless, these CNN-based algorithms show limitations in capturing long-range spatial dependencies and modeling the non-local self-similarity. Besides, the sparsity property of HSI representations is not well addressed, posing a low-efficiency problem to HSI reconstruction models.

2 Vision Transformer

Transformer is proposed for machine translation in NLP. Recently, it has gained much popularity in computer vision because of its superiority in modeling long-range interactions between spatial regions. Vision Transformer has been widely applied in image classification , object detection , semantic segmentation , human pose estimation , and so on. Besides high-level vision, Transformer has also been used in image restoration . For example, Cai et al. propose the first Transformer-based model MST for HSI reconstruction. MST treats spectral maps as tokens and calculates the self-attention along the spectral dimension. In addition, Wang et al. propose a U-shaped Transformer, named UFormer, built by the basic blocks of Swin Transformer for natural image restoration. However, existing Transformers densely sample tokens, some of which corresponding to the regions with limited information, and calculate MSA between some tokens that are unrelated in content. How to embed HSI spatial sparsity into Transformer to boost the model efficiency still remains under-studied. Our work aims to fill this research gap.

Mathematical Model of CASSI

Method

The overall framework of our coarse-to-fine sparse Transformer (CST) is shown in Fig. 3. CST consists of two key components, i.e., spectra-aware screening mechanism (SASM) for coarse patch selecting and spectra-aggregation hashing multi-head self-attention (SAH-MSA) for fine pixel clustering. Fig. 3 (a) depicts SASM and the network architecture of CST. Fig. 3 (b) shows the basic unit of CST, spectra-aware hashing attention block (SAHAB). Fig. 3 (c) illustrates our SAH-MSA, which is the most important component of SAHAB.

In our implementation, we set the basic channel $C$ = $N_{\lambda}$ = 28 to store the HSI information and change the combination ($N_{1}$,$N_{2}$,$N_{3}$) in Fig. 3 (a) to establish our CST families with small, medium, and large model sizes and computational complexities. They are CST-S (1,1,2), CST-M (2,2,2), and CST-L (2,4,6).

2 Spectra-Aware Screening Mechanism

We observe that the HSI signal exhibits high sparsity in the spatial dimension. However, the original global Transformer samples all tokens on the feature map while the window-based local Transformer samples all tokens inside every non-overlapping window. These Transformers sample many uninformative regions to calculate MSA, which degrades the model efficiency. To cope with this problem, we propose SASM for coarse patch selecting, i.e., screening out regions with dense HSI information to produce tokens. In this section, we introduce SASM in three parts, i.e., sparsity estimator, sparsity loss, and patch selection.

In this part, we detail the sparsity estimator mentioned in Sec. 4. As shown in Fig. 3 (a), the sparsity estimator adopts a U-shaped structure including a two-stage encoder, an ASSP module , and a two-stage decoder. Each stage of the encoder consists of two $conv$1$\times$1 and a strided depth-wise $conv$3$\times$3. Each stage of the decoder contains a strided $deconv$2$\times$2, two $conv$1$\times$1, and a depth-wise $conv$3$\times$3. The sparsity estimator takes the initialized feature $\mathbf{X}$ as the input to produce the shallow feature $\mathbf{X}_{0}$ and the sparsity mask $\mathbf{M}_{s}$ that localizes and screens out the informative spatial regions with HSI representations. We achieve this by minimizing our proposed sparsity loss.

2.2 Sparsity Loss.

Subsequently, our sparsity loss $\mathcal{L}_{s}$ is constructed as the mean squared error between the predicted sparsity mask $\mathbf{M}_{s}$ and the reference sparsity mask $\mathbf{M}^{*}_{s}$ as

By minimizing $\mathcal{L}_{s}$, the sparsity estimator is encouraged to detect the foreground hard-to-reconstruct regions with HSI representations. In addition, the overall training objective $\mathcal{L}$ is the weighted sum of $\mathcal{L}_{s}$ and $\mathcal{L}_{2}$ loss as

where $\mathbf{X}^{*}$ represents the ground-truth HSIs and $\lambda$ refers to the hyperparameter that controls the importance balance between $\mathcal{L}_{2}$ and $\mathcal{L}_{s}$.

2.3 Patch Selection.

3 Spectra-Aggregation Hashing Multi-head Self-Attention.

Previous Transformers calculate MSA between all the sampled tokens, some of which are even unrelated in content. This may lead to inefficient computation that lowers down the model cost-effectiveness and easily hamper convergence . The sparse coding methods assume that image signals can be represented by a sparse linear combination over dictionary signals. Inspired by this, we propose SAH-MSA for $fine~{}pixel~{}clustering$. SAH-MSA enforces a sparsity constraint on the MSA mechanism. In particular, SAH-MSA only calculates self-attention between tokens that are closely correlated in content, which addresses the limitation of previous Transformers.

where $\mathbf{B}_{i}$ represents the $i$-th $bucket$. Each $bucket$ has $m$ elements. There are $\frac{M\times M}{m}$ $buckets$ in total. With our hash clustering scheme, the closely content-correlated tokens are grouped into the same $bucket$. Therefore, the model can reduce the computational burden between content-unrelated elements by only applying the MSA operation to the tokens within the same $bucket$. More specifically, for a $query$ element $\bm{q}\in\mathbf{B}_{i}$, our SAH-MSA can be formulated as

where $R$ refers to the round number and $w^{r}_{n}$ represents the weight importance of the $r$-th round in the $n$-th head, which scores the similarity between the $query$ element $\bm{q}$ and the elements belonging to $bucket$ $\mathbf{B}_{i}^{r}$. $w^{r}_{n}$ can be obtained by

Experiment

The same with TSA-Net , 28 wavelengths from 450 nm to 650 nm are derived by spectral interpolation manipulation for simulation and real experiments.

Synthetic Data. Two HSI datasets, CAVE and KAIST , are adopted for simulation experiments. CAVE contains 32 HSIs with spatial size 512$\times$512. KAIST is composed of 30 HSIs with spatial size 2704$\times$3376. Similar to , CAVE is used for training and 10 scenes from KAIST are selected for testing.

Real Data. We adopt the real HSI dataset collected by TSA-Net .

Evaluation Metrics. We use peak signal-to-noise ratio (PSNR) and structural similarity (SSIM) as metrics to evaluate HSI reconstruction methods.

Implementation Details. Our CST models are implemented by Pytorch. They are trained with Adam optimizer ($\beta_{1}$ = 0.9 and $\beta_{2}$ = 0.999) using Cosine Annealing scheme for 500 epochs. The learning rate is initially set to 4$\times$10-4. In simulation experiments, patches at the spatial size of 256$\times$256 are randomly cropped from the 3D HSI cubes with 28 channels as training samples. For real HSI reconstruction, we set the spatial size of patches to 660$\times$660 with the same size of the real physical mask. We set the shifting step $d$ in the dispersion to 2. After the mask modulation, the image cube is sheared with an accumulative two-pixel step. Hence, the spatial sizes of measurements are 256$\times$310 and 660$\times$714 in simulation and real experiments. The batch size is set to 5. $r$ and $m$ in Eq. (7) and (8) are set to 1 and 64. The training data is augmented with random rotation and flipping. All CST models are trained and tested on a single RTX 3090 GPU.

2 Quantitative Results

We compare the Params, FLOPS, PSNR, and SSIM of our CST and other SOTA methods, including three model-based methods (TwIST , GAP-TV , and DeSCI ), six CNN-based methods ($\lambda$-net , HSSP , DNU , PnP-DIP-HSI , TSA-Net , and DGSMP ), and a recent Transformer-based method (MST ). For fairness, we test all these algorithms with the same settings as . The results on 10 simulation scenes are reported in Tab. 1. As can be seen: (i) When we set the sparsity ratio to 0, our best model CST-L∗ achieves very impressive results, i.e., 36.12 dB in PSNR and 0.957 in SSIM, showing the effectiveness of our method. (ii) Our CST families significantly outperform other SOTA algorithms while requiring cheaper computational costs. Particularly, when compared to the recent best Transformer-based method MST, our CST-S, CST-M, and CST-L achieve 0.45, 0.37, and 0.67 dB improvements while costing 1.29G, 1.16G, and 0.34G less FLOPS than MST-S, MST-M, and MST-L as shown in Fig. 5. When compared to CNN-based methods, our CST exhibits extreme efficiency advantages. For instance, CST-L outperforms DGSMP, TSA-Net, and $\lambda$-Net by 3.22, 4.39, and 7.32 dB while costing 79.8% (3.00 / 3.76), 6.8%, 4.8% Params and 4.3% (27.81 / 646.65), 25.3%, 23.6% FLOPS. Surprisingly, even our smallest model CST-S surpasses DGSMP, TSA-Net, and $\lambda$-Net by 2.08, 3.25, and 6.18 dB while requiring 31.9%, 2.7%, 1.9% Params and 1.8%, 10.6%, 9.9% FLOPS. These results demonstrate the cost-effectiveness superiority of our CST. This is mainly because CST embeds the HSI sparsity into the learning-based model, which reduces the inefficient computation of less informative dark regions and self-attention between content-unrelated tokens.

3 Qualitative Results

Fig. 4 compares the restored simulation HSIs of our CST-L and seven SOTA algorithms on $Scene$ 2 with 4 out of 28 spectral channels. Please zoom in for better visualization. It can be observed from the reconstructed HSIs (right) and the zoomed-in patches in the yellow boxes that our CST is effective in producing perceptually pleasant images with more vivid sharp edge details while maintaining the spatial smoothness of the homogeneous regions without introducing artifacts. In contrast, other methods fail to restore fine-grained details. They either achieve over-smooth results sacrificing structural contents and high-frequency details, or generate blotchy textures and chromatic artifacts. Besides, Fig. 4 depicts the spectral density curves (bottom-left) corresponding to the selected region of the green box in the RGB image (top-left). Our curve achieves the highest correlation coefficient with the ground-truth curve. This evidence clearly demonstrates the spectral-dimension consistency reconstruction effectiveness of our proposed CST.

3.2 Real HSI Restoration.

We also evaluate our CST in real HSI reconstruction. Following the setting of , we re-train our CST-L with all samples of the KAIST and CAVE datasets. To simulate real CASSI, 11-bit shot noise is injected into the measurement during the training procedure. The reconstructed HSI comparisons are depicted in Fig. 6. Our CST-L shows significant advantages in fine-grained content restoration and real noise removal. These results verify the robustness, reliability, and generalization ability of our method.

4 Ablation Study

We adopt the simulation HSI datasets to conduct ablation studies. The baseline model is derived by removing our SAH-MSA and SASM from CST-M.

We firstly perform a break-down ablation to investigate the effect of each component and their interactions. The results are listed in Tab. 2a. The baseline model yields 32.57 dB in PSNR and 0.906 in SSIM. When SAH-MSA is applied, the performance gains by 2.96 dB in PSNR and 0.042 in SSIM, showing its significant contribution. When we continue to exploit SASM, the computational cost dramatically declines by 31.3% (7.69 / 24.60) while the performance only degrades by 0.6 % in PSNR and 0.1% in SSIM. This evidence suggests that our SASM can reduce the computational burden while sacrificing minimal reconstruction performance, thus increasing the model efficiency.

4.2 Sparsity Scheme Comparison.

We conduct ablation to study the effects of sparsity schemes including: (i) random sparsity, i.e., the patches to be calculated are randomly selected, (ii) uniform sparsity, i.e., the patches to be calculated are uniformly distributed, and (iii) our SASM. The results are listed in Tab. 2b. Our SASM yields the best results and drastically outperforms other schemes (over 0.9 dB). Additionally, we conduct visual analysis of the sparsity mask generated by the three sparsity schemes. As depicted in Fig. 7, the sparsity mask produced by our SASM generates more complete and accurate responses to the informative regions with HSI information. In contrast, both random and uniform sparsity schemes are not aware of HSI signals and rigidly pick the preset positions. These results demonstrate the superiority of our SASM in perceiving spatially sparse HSI signals and locating regions with dense HSI representations.

4.3 Self-Attention Mechanism Comparison.

We compare our SAH-MSA with other self-attention mechanisms. The results are reported in Tab. 2c. The baseline yields 32.57 dB with 0.51 M Params and 6.40 G FLOPS. We respectively apply global MSA (G-MSA) , local window-based MSA (W-MSA) , Swin-MSA , spectral-wise MSA (S-MSA) , and SAH-MSA. The model gains by 2.47, 2.45, 2.55, 2.64, and 2.96 dB while adding 29.18, 18.58, 18.58, 18.34, and 18.20 G FLOPS and 1.34, 1.34, 1.34, 1.15, and 0.85 M Params. Our SAH-MSA yields the most significant improvement but requires the cheapest FLOPS and Params. Please note that we downscale the input feature of G-MSA into $\frac{1}{4}$ size to avoid memory bottlenecks. This evidence shows the cost-effectiveness advantage of SAH-MSA, which is mainly because SAH-MSA applies MSA calculation between tokens that are closely related in content within each $bucket$ while cutting down the burden of computation between content-uncorrelated elements.

4.4 Clustering Scope.

4.5 Parameter Analysis.

We adopt CST-M to conduct parameter analysis of sparsity rate $\sigma$, round number $R$ in Eq. (11), patch size $M$, and loss weight $\lambda$ in Eq. (6) as shown in Fig 8, where the vertical axis is PSNR and the circle radius is FLOPS. As can be observed: (i) When increasing $\sigma$, the computational cost declines but the performance is sacrificed. When $\sigma$ is larger than $50\%$, the performance degrades dramatically. (ii) When changing $R$ from 1 to 6, the reconstruction quality increases. Nonetheless, when $R\geq$ 2, further increasing $R$ does not lead to a significant improvement. (iii) The two maximums are achieved when $M$ = 16 and $\lambda$ = 2, respectively, without costing too much FLOPS. Since our goal is not to pursue the best results with heavy computational burden sacrificing the model efficiency but to yield a better trade-off between performance and computational cost, we finally set $\sigma$ = 0.5, $R$ = 2, $M$ = 16, and $\lambda$ = 2.

Conclusion

In this paper, we investigate a critical problem in HSI reconstruction, i.e., how to embed HSI sparsity into learning-based algorithms. To this end, we propose a novel Transformer-based method, named CST, for HSI restoration. CST firstly exploits SASM to detect informative regions with HSI representations. Then the detected patches are fed into our SAH-MSA to cluster spatially scattered tokens with closely correlated contents for calculating MSA. Extensive quantitative and qualitative experiments demonstrate that our CST significantly outperforms other SOTA methods while requiring cheaper computational costs. Additionally, our CST yields more visually pleasing results with more fine-grained details and structural contents than existing algorithms in real-world HSI reconstruction.

Acknowledgements: This work is partially supported by the NSFC fund (61831 014), the Shenzhen Science and Technology Project under Grant (JSGG20210802 153150005, CJGJZD20200617102601004), and the Westlake Foundation (2021B1 501-2). Xin Yuan would like to thank the funding from Lochn Optics.

References