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

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

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.

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

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.

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/ .

cd DWMT-main/simulation/test/

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

• 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.

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/

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.

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}
}