Uncertainty-Driven Spectral Compressive Imaging with Spatial-Frequency Transformer

Recently, learning-based Hyperspectral image (HSI) reconstruction methods have demonstrated promising performance and dominated the mainstream research direction. However, existing learning-based methods still have two issues. 1) Unable to consider both the spatial sparsity and the inter-spectral similarity prior of HSI. 2) Treat all regions equally, ignoring that texture-rich regions and edge regions are more difficult to reconstruct than smooth regions. To address these issues, we propose an uncertainty-driven HSI reconstruction method termed Specformer. Specifically, we first introduce a frequency-wise self-attention (FWSA) and combine it with spatial-wise local-window self-attention (LWSA) with a parallel design to form a Spatial-Frequency (SF) block. LWSA can guide the network to focus on the regions with dense spectral information, and FWSA can capture the inter-spectral similarity. Parallel design helps the network to model cross-window connections, expand its receptive fields while maintaining linear complexity. We use SF-block as the main building block in a multi-scale U-shape network to form our Specformer. In addition, we introduce an uncertainty-driven self-adaptive loss function, which can reinforce the network's attention to challenging regions with rich textures and edges. Comprehensive experiments show that our Specformer significantly outperforms state-of-the-art methods on simulation and real HSI datasets while requiring cheaper computational and memory costs. The code will be publicly available.