Mask-guided Spectral-wise Transformer for Efficient Hyperspectral Image Reconstruction (2111.07910v2)

Abstract: Hyperspectral image (HSI) reconstruction aims to recover the 3D spatial-spectral signal from a 2D measurement in the coded aperture snapshot spectral imaging (CASSI) system. The HSI representations are highly similar and correlated across the spectral dimension. Modeling the inter-spectra interactions is beneficial for HSI reconstruction. However, existing CNN-based methods show limitations in capturing spectral-wise similarity and long-range dependencies. Besides, the HSI information is modulated by a coded aperture (physical mask) in CASSI. Nonetheless, current algorithms have not fully explored the guidance effect of the mask for HSI restoration. In this paper, we propose a novel framework, Mask-guided Spectral-wise Transformer (MST), for HSI reconstruction. Specifically, we present a Spectral-wise Multi-head Self-Attention (S-MSA) that treats each spectral feature as a token and calculates self-attention along the spectral dimension. In addition, we customize a Mask-guided Mechanism (MM) that directs S-MSA to pay attention to spatial regions with high-fidelity spectral representations. Extensive experiments show that our MST significantly outperforms state-of-the-art (SOTA) methods on simulation and real HSI datasets while requiring dramatically cheaper computational and memory costs. Code and pre-trained models are available at https://github.com/caiyuanhao1998/MST/

• The paper introduces a mask-guided spectral transformer that leverages spectral-wise self-attention to improve hyperspectral image reconstruction.
• It integrates coded aperture mask guidance to enhance spatial fidelity and accurately model inter-spectral dependencies.
• Experimental results on benchmark datasets demonstrate up to 2.55 dB PSNR improvement and reduced computational load compared to state-of-the-art methods.

Overview of Mask-guided Spectral-wise Transformer for Efficient Hyperspectral Image Reconstruction

The paper introduces a transformative approach for reconstructing hyperspectral images (HSIs) using a novel technique named Mask-guided Spectral-wise Transformer (MST). HSIs are critical in various applications due to their ability to capture rich spectral information across different wavelengths. However, efficiently reconstructing these images from coded measurements poses significant challenges, particularly in modeling inter-spectral interactions and leveraging spatial dependencies.

Key Contributions

1. Spectral-wise Multi-head Self-Attention (S-MSA): The authors introduce S-MSA, which treats each spectral feature as a token and focuses on capturing dependencies across the spectral dimension rather than the spatial dimension. This approach effectively models the natural spectral correlation found in HSIs.
2. Mask-guided Mechanism (MM): A pivotal component of the proposed framework is MM, which leverages the coded aperture used in the CASSI (Coded Aperture Snapshot Spectral Imaging) system to guide attention mechanisms. This guidance ensures that the neural network emphasizes high-fidelity spectral regions during reconstruction.
3. Computational Efficiency: The MST framework has been demonstrated to outperform state-of-the-art (SOTA) approaches both qualitatively and quantitatively on numerous benchmark datasets while significantly reducing computational and memory requirements.

Experimental Findings and Implications

Through extensive experimentation on the CAVE and KAIST datasets, the MST models — particularly MST-L, MST-M, and MST-S — yield superior performance in terms of PSNR and SSIM metrics, compared to traditional CNN-based methods like TSA-Net and DGSMP. For example, MST-L surpasses the performance of existing methods by up to 2.55 dB in PSNR, demonstrating its efficacy in reconstructing high-fidelity HSIs while maintaining computational efficiency with lower parameters and FLOPS.

This remarkable performance is achieved by:

• Modeling Long-range Spectral Dependencies: Unlike traditional CNNs, which struggle with capturing spectral similarities and long-range dependencies, the proposed S-MSA effectively models these aspects, resulting in more accurate HSI reconstructions.
• Integrating Physical Mask Information: The use of MM allows the model to incorporate valuable spatial fidelity information during the learning process, which is often overlooked in earlier methods. This integration leads to enhanced spatial attention and consequently better spectral estimation.

Theoretical and Practical Implications

The MST framework presents a paradigm shift in how HSIs are reconstructed, emphasizing the spectral dimension rather than conventional spatial-focused techniques. This perspective may pave the way for future AI models that require efficient representation and processing of multi-dimensional data, extending beyond hyperspectral imaging to other domains involving volumetric or time-series data analysis.

The implementation also highlights the role of domain-specific insights, such as leveraging physical mask information, in enhancing model performance. This approach suggests that similar methodologies could be beneficially adapted in areas like remote sensing, medical imaging, and environmental monitoring where rich spectral data is abundant.

Future Directions

Potential future research avenues include:

• Scaling and Generalization: Exploring the scalability of MST models for larger and more diverse datasets could provide insights into their applicability in other imaging domains.
• Real-time Reconstruction: Enhancing the efficiency of MST models to enable real-time HSI reconstruction would be valuable in dynamic scenarios, such as real-time monitoring and surveillance.
• Integration with Other Modalities: Extending the framework to integrate with other imaging or data modalities, such as LiDAR or radar, can potentially unlock new capabilities in multi-sensor data fusion applications.

The paper highlights the transformative potential of spectral-based attention mechanisms in hyperspectral image processing, providing a robust and efficient solution that bridges the gap between computational feasibility and high reconstruction fidelity.