CaiYuanhao/MST_pp

MST++: Multi-stage Spectral-wise Transformer for Efficient Spectral Reconstruction

CVPR Workshops 2022 · Winner of the NTIRE 2022 Spectral Reconstruction Challenge

Yuanhao Cai, Jing Lin, Zudi Lin, Haoqian Wang, Yulun Zhang, Hanspeter Pfister, Radu Timofte, and Luc Van Gool

The first two authors contribute equally.

Code and project documentation • Paper on Hugging Face

Model description

MST++ reconstructs hyperspectral images from RGB inputs. It uses attention across spectral channels to capture spectral relationships, with U-shaped Single-stage Spectral-wise Transformers extracting features at multiple resolutions. Cascading these stages progressively refines the reconstruction.

This repository provides pretrained PyTorch weights for MST++ and 10 comparison methods from the official spectral reconstruction toolbox. Use each checkpoint with its corresponding architecture in the code repository.

NTIRE awards

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

The weights hosted here are for RGB-to-hyperspectral reconstruction. The NTIRE 2024 low-light enhancement weights are available in CaiYuanhao/Retinexformer.

Visual examples

480 nm	520 nm	580 nm	660 nm

Network architecture

Available checkpoints

Model	Checkpoint	--method
MST++	mst_plus_plus.pth	mst_plus_plus
MST-L	mst.pth	mst
HDNet	hdnet.pth	hdnet
Restormer	restormer.pth	restormer
MPRNet	mprnet.pth	mprnet
HINet	hinet.pth	hinet
MIRNet	mirnet.pth	mirnet
AWAN	awan.pth	awan
HRNet	hrnet.pth	hrnet
HSCNN+	hscnn_plus.pth	hscnn_plus
EDSR	edsr.pth	edsr

Getting started

1. Install the project

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

The official inference scripts require an NVIDIA GPU, CUDA, and a compatible PyTorch installation. See the project README for environment and dataset preparation.

2. Download model weights

Run from the MST-plus-plus repository root:

pip install -U huggingface_hub
hf download CaiYuanhao/MST_pp --include "*.pth" --local-dir predict_code/model_zoo

To download only MST++:

hf download CaiYuanhao/MST_pp mst_plus_plus.pth --local-dir predict_code/model_zoo

3. Reconstruct an RGB image

cd predict_code
python test.py \
  --rgb_path ./demo/ARAD_1K_0912.jpg \
  --method mst_plus_plus \
  --pretrained_model_path ./model_zoo/mst_plus_plus.pth \
  --outf ./exp/mst_plus_plus/ \
  --gpu_id 0

Replace --rgb_path with your RGB image path. The reconstructed hyperspectral image is saved in MATLAB .mat format under predict_code/exp/mst_plus_plus/.

4. Run the validation or test scripts

Prepare the dataset using the official instructions. For validation, run the following from the repository root:

hf download CaiYuanhao/MST_pp --include "*.pth" --local-dir test_develop_code/model_zoo
cd test_develop_code
python test.py \
  --data_root ../dataset/ \
  --method mst_plus_plus \
  --pretrained_model_path ./model_zoo/mst_plus_plus.pth \
  --outf ./exp/mst_plus_plus/ \
  --gpu_id 0

For the challenge test set, download weights into test_challenge_code/model_zoo, then run the same testing command from test_challenge_code/. See the GitHub documentation for dataset splits, other methods, training commands, and visualization tools.

Citation

# MST
@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}
}


# MST++
@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}
}


# HDNet
@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}
}