DSMT: Dual-Stage Multiscale Transformer for Hyperspectral Snapshot Compressive Imaging

This repo is the implementation of DSMT: Dual-Stage Multiscale Transformer for Hyperspectral Snapshot Compressive Imaging (TIP2025).

Acknowledgements

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

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
DSMT (ours)	12.40	49.94	36.92	0.966

Qualitative Results

Download results of DSMT (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:

|--DSMT-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 DSMT-main/simulation/train/

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

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

Testing

cd DSMT-main/simulation/test/

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

Visualization

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

• Generate the RGB images of the reconstructed HSIs.

Real Experiement:

Training

cd DSMT-main/real/train/

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

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

Testing

cd DSMT-main/real/test/

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

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

Visualization

• Put the reconstruted HSI in --pretrained_model_path ./checkpoints/dsmt.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:

@article{luo2025dsmt,
  title={DSMT: Dual-Stage Multiscale Transformer for Hyperspectral Snapshot Compressive Imaging},
  author={Luo, Fulin and Chen, Xi and Guo, Tan and Gong, Xiuwen and Zhang, Lefei and Zhu, Ce},
  journal={IEEE Transactions on Image Processing},
  year={2025},
  volume={34},
  number={},
  pages={2473-2486}
}