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

Network architecture

Available checkpoints

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

The root config.json lists this repository's checkpoints, supporting files, and implementation. Include it when downloading selected weights so those downloads can be counted by the Hugging Face Hub. Use the linked implementation and its matching runtime configuration to load the checkpoints.

Run from the MST-plus-plus repository root:

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

To download only MST++:

hf download CaiYuanhao/MST_pp config.json 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 "config.json" "*.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}
}