MST: A Toolbox for Spectral Compressive Imaging

Code and project documentation

Overview

This repository provides pretrained PyTorch checkpoints for spectral compressive imaging reconstruction, including MST, CST, DAUHST, BiSRNet, MST++, HDNet, and other learning-based baselines supported by the MST toolbox.

The toolbox reconstructs hyperspectral image cubes from coded snapshot measurements using the corresponding sensing mask and measurement model. The simulation implementation reconstructs 28 spectral bands. This task differs from reconstructing spectra directly from ordinary RGB photographs.

MST — Mask-guided Spectral-wise Transformer for Efficient Hyperspectral Image Reconstruction (CVPR 2022) models dependencies across spectral channels and incorporates sensing-mask information to guide reconstruction.

MST authors: Yuanhao Cai, Jing Lin, Xiaowan Hu, Haoqian Wang, Xin Yuan, Yulun Zhang, Radu Timofte, and Luc Van Gool.

Related papers

• MST: Mask-guided Spectral-wise Transformer for Efficient Hyperspectral Image Reconstruction — CVPR 2022.
• CST: Coarse-to-Fine Sparse Transformer for Hyperspectral Image Reconstruction — ECCV 2022.
• DAUHST: Degradation-Aware Unfolding Half-Shuffle Transformer for Spectral Compressive Imaging — NeurIPS 2022.
• BiSRNet: Binarized Spectral Compressive Imaging — NeurIPS 2023.
• MST++: Multi-stage Spectral-wise Transformer for Efficient Spectral Reconstruction — CVPR Workshops 2022.
• HDNet: High-resolution Dual-domain Learning for Spectral Compressive Imaging — CVPR 2022.

NTIRE awards

• 🏆 NTIRE 2022 — Winner: MST++ won the Spectral Recovery from RGB Images Challenge. See the MST++ project.
• 🏆 NTIRE 2024 — First Runner-Up: Retinexformer and MST++ achieved second place in the Low-Light Image Enhancement Challenge.

The awards concern the associated challenge solutions. This repository hosts the spectral compressive imaging model zoo; the NTIRE 2024 low-light enhancement weights are available in CaiYuanhao/Retinexformer.

Visual examples

Scene 2	Scene 3	Scene 4	Scene 7

See the GitHub repository for reconstruction comparisons and visualization tools.

Available checkpoints

The Hugging Face repository contains 20 checkpoints across 12 model families:

Model family	Checkpoint paths
ADMM-Net	admm_net/admm_net.pth
BIRNAT	birnat/birnat.pth
BiSRNet	bisrnet/bisrnet.pth
CST	cst/cst_l.pth cst/cst_l_plus.pth cst/cst_m.pth cst/cst_s.pth
DAUHST	dauhst/dauhst_2stg.pth dauhst/dauhst_3stg.pth dauhst/dauhst_5stg.pth dauhst/dauhst_9stg.pth
DGSMP	dgsmp/dgsmp.pth
GAP-Net	gap_net/gap_net.pth
HDNet	hdnet/hdnet.pth
λ-Net	lambda_net/lambda_net.pth
MST	mst/mst_l.pth mst/mst_m.pth mst/mst_s.pth
MST++	mst_plus_plus/mst_plus_plus.pth
TSA-Net	tsa_net/tsa_net.pth

Each checkpoint requires its matching architecture and method setting in the toolbox.

Getting started

1. Install the toolbox

git clone https://github.com/caiyuanhao1998/MST.git
cd MST
pip install -r requirements.txt

The official testing implementation requires an NVIDIA GPU with CUDA and a compatible PyTorch installation. See the environment instructions.

2. Download model weights

Run from the MST repository root to download the weights into the simulation model zoo while preserving the expected directory structure:

pip install -U huggingface_hub
hf download CaiYuanhao/MST --include "*.pth" --local-dir simulation/test_code/model_zoo

To download only MST-S:

hf download CaiYuanhao/MST mst/mst_s.pth --local-dir simulation/test_code/model_zoo

For example, the MST-S checkpoint will be saved as simulation/test_code/model_zoo/mst/mst_s.pth.

3. Prepare the data

Follow the project's dataset preparation instructions to place the hyperspectral test scenes and sensing mask under datasets/TSA_simu_data/:

datasets/
└── TSA_simu_data/
    ├── mask.mat
    └── Truth/
        ├── scene01.mat
        ├── scene02.mat
        └── ...

4. Run simulation reconstruction

From the MST repository root:

cd simulation/test_code

# MST-S
python test.py --template mst_s --method mst_s \
  --pretrained_model_path ./model_zoo/mst/mst_s.pth \
  --outf ./exp/mst_s/

# MST++
python test.py --template mst_plus_plus --method mst_plus_plus \
  --pretrained_model_path ./model_zoo/mst_plus_plus/mst_plus_plus.pth \
  --outf ./exp/mst_plus_plus/

# DAUHST-9stg
python test.py --template dauhst_9stg --method dauhst_9stg \
  --pretrained_model_path ./model_zoo/dauhst/dauhst_9stg.pth \
  --outf ./exp/dauhst_9stg/

Reconstructed hyperspectral images are saved as Test_result.mat under the corresponding output directory. See the official documentation for other methods, real-measurement experiments, MATLAB quality metrics, and visualization scripts.

Citation

Please cite the paper corresponding to the method you use:

@inproceedings{mst,
  title={Mask-guided Spectral-wise Transformer for Efficient Hyperspectral Image Reconstruction},
  author={Yuanhao Cai and Jing Lin and Xiaowan Hu and Haoqian Wang and Xin Yuan and Yulun Zhang and Radu Timofte and Luc Van Gool},
  booktitle={CVPR},
  year={2022}
}

@inproceedings{cst,
  title={Coarse-to-Fine Sparse Transformer for Hyperspectral Image Reconstruction},
  author={Yuanhao Cai and Jing Lin and Xiaowan Hu and Haoqian Wang and Xin Yuan and Yulun Zhang and Radu Timofte and Luc Van Gool},
  booktitle={ECCV},
  year={2022}
}

@inproceedings{dauhst,
  title={Degradation-Aware Unfolding Half-Shuffle Transformer for Spectral Compressive Imaging},
  author={Yuanhao Cai and Jing Lin and Haoqian Wang and Xin Yuan and Henghui Ding and Yulun Zhang and Radu Timofte and Luc Van Gool},
  booktitle={NeurIPS}, 
  year={2022}
}

@inproceedings{bisci,
  title={Binarized Spectral Compressive Imaging},
  author={Yuanhao Cai and Yuxin Zheng and Jing Lin and Xin Yuan and Yulun Zhang and Haoqian Wang},
  booktitle={NeurIPS},
  year={2023}
}

@inproceedings{mst_pp,
  title={MST++: Multi-stage Spectral-wise Transformer for Efficient Spectral Reconstruction},
  author={Yuanhao Cai and Jing Lin and Zudi Lin and Haoqian Wang and Yulun Zhang and Hanspeter Pfister and Radu Timofte and Luc Van Gool},
  booktitle={CVPRW},
  year={2022}
}

@inproceedings{hdnet,
  title={HDNet: High-resolution Dual-domain Learning for Spectral Compressive Imaging},
  author={Xiaowan Hu and Yuanhao Cai and Jing Lin and  Haoqian Wang and Xin Yuan and Yulun Zhang and Radu Timofte and Luc Van Gool},
  booktitle={CVPR},
  year={2022}
}