DWMT for Hyperspectral Snapshot Compressive Imaging

This repo is the implementation of Dual-Window Multiscale Transformer for Hyperspectral Snapshot Compressive Imaging (AAAI2024).

Acknowledgements

Our code is heavily borrows from CST (ECCV 2022) and DAUHST (NeurIPS 2022). Thanks to their generous open source efforts!

Abstract

Coded aperture snapshot spectral imaging (CASSI) system is an effective manner for hyperspectral snapshot compressive imaging. The core issue of CASSI is to solve the inverse problem for the reconstruction of hyperspectral image (HSI). In recent years, Transformer-based methods achieve promising performance in HSI reconstruction. However, capturing both long-range dependencies and local information while ensuring reasonable computational costs remains a challenging problem. In this paper, we propose a Transformer-based HSI reconstruction method called dual-window multiscale Transformer (DWMT), which is a coarse-to-fine process, reconstructing the global properties of HSI with the long-range dependencies. In our method, we propose a novel U-Net architecture using a dual-branch encoder to refine pixel information and full-scale skip connections to fuse different features, enhancing the extraction of fine-grained features. Meanwhile, we design a novel self-attention mechanism called dual-window multiscale multi-head self-attention (DWM-MSA), which utilizes two different-sized windows to compute self-attention, which can capture the long-range dependencies in a local region at different scales to improve the reconstruction performance. We also propose a novel position embedding method for Transformer, named con-abs position embedding (CAPE), which effectively enhances positional information of the HSIs. Extensive experiments on both the simulated and the real data are conducted to demonstrate the superior performance, stability, and generalization ability of our DWMT.

Comparison on Simulation Dataset

The performance are reported on 10 scenes of the KAIST dataset. The test size of FLOPS is 256 x 256.

Quantitative Results

Method	Params (M)	FLOPS (G)	PSNR	SSIM
TSA-Net	44.25	110.06	31.46	0.894
HDNet	2.37	154.76	34.97	0.943
MST	2.03	28.15	35.18	0.948
CST	3.00	40.10	36.12	0.957
DWMT (ours)	14.48	46.71	36.82	0.964

Qualitative Results

Download results of DWMT (Google Drive).

Usage

Prepare Dataset

Download cave_1024_28 (Baidu Disk, code: fo0q | One Drive), CAVE_512_28 (Baidu Disk, code: ixoe | One Drive), KAIST_CVPR2021 (Baidu Disk, code: 5mmn | One Drive), TSA_simu_data (Baidu Disk, code: efu8 | One Drive), TSA_real_data (Baidu Disk, code: eaqe | One Drive), and then put them into the corresponding folders of datasets/ and recollect them as the following form:

|--DWMT-main
    |--real
    	|-- test
    	|-- train
    |--simulation
    	|-- test
    	|-- train
    |--visualization
    |--datasets
        |--cave_1024_28
            |--scene1.mat
            |--scene2.mat
            ：  
            |--scene205.mat
        |--CAVE_512_28
            |--scene1.mat
            |--scene2.mat
            ：  
            |--scene30.mat
        |--KAIST_CVPR2021  
            |--1.mat
            |--2.mat
            ： 
            |--30.mat
        |--TSA_simu_data  
            |--mask.mat   
            |--Truth
                |--scene01.mat
                |--scene02.mat
                ： 
                |--scene10.mat
        |--TSA_real_data  
            |--mask.mat   
            |--Measurements
                |--scene1.mat
                |--scene2.mat
                ： 
                |--scene5.mat

Following CST and DAUHST, we use the CAVE dataset (cave_1024_28) as the simulation training set. Both the CAVE (CAVE_512_28) and KAIST (KAIST_CVPR2021) datasets are used as the real training set.

Simulation Experiement:

Training

cd DWMT-main/simulation/train/

python train.py --template dwmt --outf ./exp/dwmt/ --method dwmt

The training log, trained model, and reconstrcuted HSI will be available in DWMT-main/simulation/train/exp/ .

Testing

cd DWMT-main/simulation/test/

python test.py --template dwmt --outf ./exp/dwmt/ --method dwmt --pretrained_model_path ./checkpoints/dwmt.pth

Visualization

• Put the reconstruted HSI in DWMT-main/visualization/simulation_results/results and rename it as method.mat.

• Generate the RGB images of the reconstructed HSIs.

Real Experiement:

Training

cd DWMT-main/real/train/

python train.py --template dwmt --outf ./exp/dwmt/ --method dwmt

The training log, trained model, and reconstrcuted HSI will be available in DWMT-main/real/train/exp/

Testing

cd DWMT-main/real/test/

python test.py --template dwmt --outf ./exp/dwmt/ --method dwmt --pretrained_model_path ./checkpoints/dwmt.pth

• The reconstrcuted HSI will be output into DWMT-main/real/test/exp/

###　Visualization

• Put the reconstruted HSI in --pretrained_model_path ./checkpoints/dwmt.pth/visualization/real_results/results and rename it as method.mat, e.g., mst_plus_plus.mat.

• Generate the RGB images of the reconstructed HSI.

Citation

If this code helps you, please consider citing our work:

@inproceedings{luo2024dual,
  title={Dual-Window Multiscale Transformer for Hyperspectral Snapshot Compressive Imaging},
  author={Luo, Fulin and Chen, Xi and Gong, Xiuwen and Wu, Weiwen and Guo, Tan},
  booktitle={Proceedings of the AAAI Conference on Artificial Intelligence},
  volume={38},
  number={4},
  pages={3972--3980},
  year={2024}
}