https://github.com/bryanbocao/mst

A toolbox for spectral compressive imaging reconstruction including our MST (CVPR 2022), CST (ECCV 2022), DAUHST (NeurIPS 2022), HDNet (CVPR 2022), MST++ (CVPRW 2022), etc.

https://github.com/bryanbocao/mst

Science Score: 23.0%

This score indicates how likely this project is to be science-related based on various indicators:

  • CITATION.cff file
  • codemeta.json file
  • .zenodo.json file
  • DOI references
    Found 2 DOI reference(s) in README
  • Academic publication links
    Links to: arxiv.org, springer.com, ieee.org
  • Academic email domains
  • Institutional organization owner
  • JOSS paper metadata
  • Scientific vocabulary similarity
    Low similarity (7.0%) to scientific vocabulary
Last synced: 10 months ago · JSON representation

Repository

A toolbox for spectral compressive imaging reconstruction including our MST (CVPR 2022), CST (ECCV 2022), DAUHST (NeurIPS 2022), HDNet (CVPR 2022), MST++ (CVPRW 2022), etc.

Basic Info
  • Host: GitHub
  • Owner: bryanbocao
  • Language: Python
  • Default Branch: main
  • Homepage:
  • Size: 5.21 MB
Statistics
  • Stars: 0
  • Watchers: 0
  • Forks: 0
  • Open Issues: 0
  • Releases: 0
Fork of caiyuanhao1998/MST
Created almost 3 years ago · Last pushed over 2 years ago
Metadata Files
Readme

README.md

A Toolbox for Spectral Compressive Imaging

winner zhihu zhihu zhihu

Authors

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

Papers

Awards

ntire

News

  • 2023.02.26 : We release the RGB images of five real scenes and ten simulation scenes. Please feel free to check and use them. 🌟
  • 2022.11.02 : We have provided more visual results of state-of-the-art methods and the function to evaluate the parameters and computational complexity of models. Please feel free to check and use them. :high_brightness:
  • 2022.10.23 : Code, models, and recontructed HSI results of DAUHST have been released. 🔥
  • 2022.09.15 : Our DAUHST has been accepted by NeurIPS 2022, code and models are coming soon. :rocket:
  • 2022.07.20 : Code, models, and recontructed HSI results of CST have been released. 🔥
  • 2022.07.04 : Our paper CST has been accepted by ECCV 2022, code and models are coming soon. :rocket:
  • 2022.06.14 : Code and models of MST and MST++ have been released. This repo supports 11 learning-based methods to serve as toolbox for Spectral Compressive Imaging. The model zoo will be enlarged. 🔥
  • 2022.05.20 : Our work DAUHST is on arxiv. :dizzy:
  • 2022.04.02 : Further work MST++ has won the NTIRE 2022 Spectral Reconstruction Challenge. :trophy:
  • 2022.03.09 : Our work CST is on arxiv. :dizzy:
  • 2022.03.02 : Our paper MST has been accepted by CVPR 2022, code and models are coming soon. :rocket:

| Scene 2 | Scene 3 | Scene 4 | Scene 7 | | :----------------------------------------------------------: | :----------------------------------------------------------: | :----------------------------------------------------------: | :----------------------------------------------------------: | | | | | |

 

1. Comparison with State-of-the-art Methods

This repo is a baseline and toolbox containing 11 learning-based algorithms for spectral compressive imaging.

Supported algorithms: * [x] [MST](https://arxiv.org/abs/2111.07910) (CVPR 2022) * [x] [CST](https://arxiv.org/abs/2203.04845) (ECCV 2022) * [x] [DAUHST](https://arxiv.org/abs/2205.10102) (NeurIPS 2022) * [x] [MST++](https://arxiv.org/abs/2111.07910) (CVPRW 2022) * [x] [HDNet](https://arxiv.org/abs/2203.02149) (CVPR 2022) * [x] [BIRNAT](https://ieeexplore.ieee.org/abstract/document/9741335/) (TPAMI 2022) * [x] [DGSMP](https://arxiv.org/abs/2103.07152) (CVPR 2021) * [x] [GAP-Net](https://arxiv.org/abs/2012.08364) (Arxiv 2020) * [x] [TSA-Net](https://link.springer.com/chapter/10.1007/978-3-030-58592-1_12) (ECCV 2020) * [x] [ADMM-Net](https://openaccess.thecvf.com/content_ICCV_2019/html/Ma_Deep_Tensor_ADMM-Net_for_Snapshot_Compressive_Imaging_ICCV_2019_paper.html) (ICCV 2019) * [x] [λ-Net](https://ieeexplore.ieee.org/document/9010044) (ICCV 2019)

We are going to enlarge our model zoo in the future.

| MST vs. SOTA | CST vs. MST | | :----------------------------------------------: | :-----------------------------------------: | | | | | MST++ vs. SOTA | DAUHST vs. SOTA | | | |

Quantitative Comparison on Simulation Dataset

| Method | Params (M) | FLOPS (G) | PSNR | SSIM | Model Zoo | Simulation Result | Real Result | | :----------------------------------------------------------: | :--------: | :-------: | :---: | :---: | :----------------------------------------------------------: | :----------------------------------------------------------: | :----------------------------------------------------------: | | λ-Net | 62.64 | 117.98 | 28.53 | 0.841 | Google Drive / Baidu Disk | Google Drive / Baidu Disk | Google Drive / Baidu Disk | | TSA-Net | 44.25 | 110.06 | 31.46 | 0.894 | Google Drive / Baidu Disk | Google Drive / Baidu Disk | Google Drive / Baidu Disk | | DGSMP | 3.76 | 646.65 | 32.63 | 0.917 | Google Drive / Baidu Disk | Google Drive / Baidu Disk | Google Drive / Baidu Disk | | GAP-Net | 4.27 | 78.58 | 33.26 | 0.917 | Google Drive / Baidu Disk | Google Drive / Baidu Disk | Google Drive / Baidu Disk | | ADMM-Net | 4.27 | 78.58 | 33.58 | 0.918 | Google Drive / Baidu Disk | Google Drive / Baidu Disk | Google Drive / Baidu Disk | | BIRNAT | 4.40 | 2122.66 | 37.58 | 0.960 | Google Drive / Baidu Disk | Google Drive / Baidu Disk | Google Drive / Baidu Disk | | HDNet | 2.37 | 154.76 | 34.97 | 0.943 | Google Drive / Baidu Disk | Google Drive / Baidu Disk | Google Drive / Baidu Disk | | MST-S | 0.93 | 12.96 | 34.26 | 0.935 | Google Drive / Baidu Disk | Google Drive / Baidu Disk | Google Drive / Baidu Disk | | MST-M | 1.50 | 18.07 | 34.94 | 0.943 | Google Drive / Baidu Disk | Google Drive / Baidu Disk | Google Drive / Baidu Disk | | MST-L | 2.03 | 28.15 | 35.18 | 0.948 | Google Drive / Baidu Disk | Google Drive / Baidu Disk | Google Drive / Baidu Disk | | MST++ | 1.33 | 19.42 | 35.99 | 0.951 | Google Drive / Baidu Disk | Google Drive / Baidu Disk | Google Drive / Baidu Disk | | CST-S | 1.20 | 11.67 | 34.71 | 0.940 | Google Drive / Baidu Disk | Google Drive / Baidu Disk | Google Drive / Baidu Disk | | CST-M | 1.36 | 16.91 | 35.31 | 0.947 | Google Drive / Baidu Disk | Google Drive / Baidu Disk | Google Drive / Baidu Disk | | CST-L | 3.00 | 27.81 | 35.85 | 0.954 | Google Drive / Baidu Disk | Google Drive / Baidu Disk | Google Drive / Baidu Disk | | CST-L-Plus | 3.00 | 40.10 | 36.12 | 0.957 | Google Drive / Baidu Disk | Google Drive / Baidu Disk | Google Drive / Baidu Disk | | DAUHST-2stg | 1.40 | 18.44 | 36.34 | 0.952 | Google Drive / Baidu Disk | Google Drive /Baidu Disk | Google Drive / Baidu Disk | | DAUHST-3stg | 2.08 | 27.17 | 37.21 | 0.959 | Google Drive / Baidu Disk | Google Drive / Baidu Disk | Google Drive / Baidu Disk | | DAUHST-5stg | 3.44 | 44.61 | 37.75 | 0.962 | Google Drive / Baidu Disk | Google Drive / Baidu Disk | Google Drive / Baidu Disk | | DAUHST-9stg | 6.15 | 79.50 | 38.36 | 0.967 | Google Drive / Baidu Disk | Google Drive / Baidu Disk | Google Drive / Baidu Disk |

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

We also provide the RGB images of five real scenes and ten simulation scenes for your convenience to draw a figure.

Note: access code for Baidu Disk is mst1

 

2. Create Environment:

  • Python 3 (Recommend to use Anaconda)

  • NVIDIA GPU + CUDA

  • Python packages:

shell pip install -r requirements.txt

 

3. Prepare Dataset:

Download cave102428 (Baidu Disk, code: fo0q | One Drive), CAVE51228 (Baidu Disk, code: ixoe | One Drive), KAISTCVPR2021 (Baidu Disk, code: 5mmn | One Drive), TSAsimudata (Baidu Disk, code: efu8 | One Drive), TSAreal_data (Baidu Disk, code: eaqe | One Drive), and then put them into the corresponding folders of datasets/ and recollect them as the following form:

shell |--MST |--real |-- test_code |-- train_code |--simulation |-- test_code |-- train_code |--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 TSA-Net and DGSMP, we use the CAVE dataset (cave102428) as the simulation training set. Both the CAVE (CAVE51228) and KAIST (KAIST_CVPR2021) datasets are used as the real training set.

 

4. Simulation Experiement:

4.1 Training

```shell cd MST/simulation/train_code/

MST_S

python train.py --template msts --outf ./exp/msts/ --method mst_s

MST_M

python train.py --template mstm --outf ./exp/mstm/ --method mst_m

MST_L

python train.py --template mstl --outf ./exp/mstl/ --method mst_l

CST_S

python train.py --template csts --outf ./exp/csts/ --method cst_s

CST_M

python train.py --template cstm --outf ./exp/cstm/ --method cst_m

CST_L

python train.py --template cstl --outf ./exp/cstl/ --method cst_l

CSTLPlus

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

GAP-Net

python train.py --template gapnet --outf ./exp/gapnet/ --method gap_net

ADMM-Net

python train.py --template admmnet --outf ./exp/admmnet/ --method admm_net

TSA-Net

python train.py --template tsanet --outf ./exp/tsanet/ --method tsa_net

HDNet

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

DGSMP

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

BIRNAT

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

MSTPlusPlus

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

λ-Net

python train.py --template lambdanet --outf ./exp/lambdanet/ --method lambda_net

DAUHST-2stg

python train.py --template dauhst2stg --outf ./exp/dauhst2stg/ --method dauhst_2stg

DAUHST-3stg

python train.py --template dauhst3stg --outf ./exp/dauhst3stg/ --method dauhst_3stg

DAUHST-5stg

python train.py --template dauhst5stg --outf ./exp/dauhst5stg/ --method dauhst_5stg

DAUHST-9stg

python train.py --template dauhst9stg --outf ./exp/dauhst9stg/ --method dauhst_9stg ```

The training log, trained model, and reconstrcuted HSI will be available in MST/simulation/train_code/exp/ .

4.2 Testing

Download the pretrained model zoo from (Google Drive / Baidu Disk, code: mst1) and place them to MST/simulation/test_code/model_zoo/

Run the following command to test the model on the simulation dataset.

```python cd MST/simulation/test_code/

MST_S

python test.py --template msts --outf ./exp/msts/ --method msts --pretrainedmodelpath ./modelzoo/mst/mst_s.pth

MST_M

python test.py --template mstm --outf ./exp/mstm/ --method mstm --pretrainedmodelpath ./modelzoo/mst/mst_m.pth

MST_L

python test.py --template mstl --outf ./exp/mstl/ --method mstl --pretrainedmodelpath ./modelzoo/mst/mst_l.pth

CST_S

python test.py --template csts --outf ./exp/csts/ --method csts --pretrainedmodelpath ./modelzoo/cst/cst_s.pth

CST_M

python test.py --template cstm --outf ./exp/cstm/ --method cstm --pretrainedmodelpath ./modelzoo/cst/cst_m.pth

CST_L

python test.py --template cstl --outf ./exp/cstl/ --method cstl --pretrainedmodelpath ./modelzoo/cst/cst_l.pth

CSTLPlus

python test.py --template cstlplus --outf ./exp/cstlplus/ --method cstlplus --pretrainedmodelpath ./modelzoo/cst/cstl_plus.pth

GAP_Net

python test.py --template gapnet --outf ./exp/gapnet/ --method gapnet --pretrainedmodelpath ./modelzoo/gapnet/gapnet.pth

ADMM_Net

python test.py --template admmnet --outf ./exp/admmnet/ --method admmnet --pretrainedmodelpath ./modelzoo/admmnet/admmnet.pth

TSA_Net

python test.py --template tsanet --outf ./exp/tsanet/ --method tsanet --pretrainedmodelpath ./modelzoo/tsanet/tsanet.pth

HDNet

python test.py --template hdnet --outf ./exp/hdnet/ --method hdnet --pretrainedmodelpath ./model_zoo/hdnet/hdnet.pth

DGSMP

python test.py --template dgsmp --outf ./exp/dgsmp/ --method dgsmp --pretrainedmodelpath ./model_zoo/dgsmp/dgsmp.pth

BIRNAT

python test.py --template birnat --outf ./exp/birnat/ --method birnat --pretrainedmodelpath ./model_zoo/birnat/birnat.pth

MSTPlusPlus

python test.py --template mstplusplus --outf ./exp/mstplusplus/ --method mstplusplus --pretrainedmodelpath ./modelzoo/mstplusplus/mstplus_plus.pth

λ-Net

python test.py --template lambdanet --outf ./exp/lambdanet/ --method lambdanet --pretrainedmodelpath ./modelzoo/lambdanet/lambdanet.pth

DAUHST-2stg

python test.py --template dauhst2stg --outf ./exp/dauhst2stg/ --method dauhst2stg --pretrainedmodelpath ./modelzoo/dauhst2stg/dauhst2stg.pth

DAUHST-3stg

python test.py --template dauhst3stg --outf ./exp/dauhst3stg/ --method dauhst3stg --pretrainedmodelpath ./modelzoo/dauhst3stg/dauhst3stg.pth

DAUHST-5stg

python test.py --template dauhst5stg --outf ./exp/dauhst5stg/ --method dauhst5stg --pretrainedmodelpath ./modelzoo/dauhst5stg/dauhst5stg.pth

DAUHST-9stg

python test.py --template dauhst9stg --outf ./exp/dauhst9stg/ --method dauhst9stg --pretrainedmodelpath ./modelzoo/dauhst9stg/dauhst9stg.pth ```

  • The reconstrcuted HSIs will be output into MST/simulation/test_code/exp/

  • Place the reconstructed results into MST/simulation/test_code/Quality_Metrics/results and

shell Run cal_quality_assessment.m

to calculate the PSNR and SSIM of the reconstructed HSIs.

  • #### Evaluating the Params and FLOPS of models

We have provided a function my_summary() in simulation/test_code/utils.py, please use this function to evaluate the parameters and computational complexity of the models, especially the Transformers as

shell from utils import my_summary my_summary(MST(), 256, 256, 28, 1)

4.3 Visualization

  • Put the reconstruted HSI in MST/visualization/simulation_results/results and rename it as method.mat, e.g., mst_s.mat.

  • Generate the RGB images of the reconstructed HSIs

shell cd MST/visualization/ Run show_simulation.m

  • Draw the spetral density lines

shell cd MST/visualization/ Run show_line.m

 

5. Real Experiement:

5.1 Training

```shell cd MST/real/train_code/

MST_S

python train.py --template msts --outf ./exp/msts/ --method mst_s

MST_M

python train.py --template mstm --outf ./exp/mstm/ --method mst_m

MST_L

python train.py --template mstl --outf ./exp/mstl/ --method mst_l

CST_S

python train.py --template csts --outf ./exp/csts/ --method cst_s

CST_M

python train.py --template cstm --outf ./exp/cstm/ --method cst_m

CST_L

python train.py --template cstl --outf ./exp/cstl/ --method cst_l

CSTLPlus

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

GAP-Net

python train.py --template gapnet --outf ./exp/gapnet/ --method gap_net

ADMM-Net

python train.py --template admmnet --outf ./exp/admmnet/ --method admm_net

TSA-Net

python train.py --template tsanet --outf ./exp/tsanet/ --method tsa_net

HDNet

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

DGSMP

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

BIRNAT

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

MSTPlusPlus

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

λ-Net

python train.py --template lambdanet --outf ./exp/lambdanet/ --method lambda_net

DAUHST-2stg

python train.py --template dauhst2stg --outf ./exp/dauhst2stg/ --method dauhst_2stg

DAUHST-3stg

python train.py --template dauhst3stg --outf ./exp/dauhst3stg/ --method dauhst_3stg

DAUHST-5stg

python train.py --template dauhst5stg --outf ./exp/dauhst5stg/ --method dauhst_5stg

DAUHST-9stg

python train.py --template dauhst9stg --outf ./exp/dauhst9stg/ --method dauhst_9stg ```

The training log, trained model, and reconstrcuted HSI will be available in MST/real/train_code/exp/

5.2 Testing

The testing scripts in real dataset should be updated >>>

```python cd MST/real/test_code/

MST_S

python test.py --template msts --outf ./exp/msts/ --method msts --pretrainedmodelpath ./modelzoo/mst/mst_s.pth

MST_M

python test.py --template mstm --outf ./exp/mstm/ --method mstm --pretrainedmodelpath ./modelzoo/mst/mst_m.pth

MST_L

python test.py --template mstl --outf ./exp/mstl/ --method mstl --pretrainedmodelpath ./modelzoo/mst/mst_l.pth

CST_S

python test.py --template csts --outf ./exp/csts/ --method csts --pretrainedmodelpath ./modelzoo/cst/cst_s.pth

CST_M

python test.py --template cstm --outf ./exp/cstm/ --method cstm --pretrainedmodelpath ./modelzoo/cst/cst_m.pth

CST_L

python test.py --template cstl --outf ./exp/cstl/ --method cstl --pretrainedmodelpath ./modelzoo/cst/cst_l.pth

CSTLPlus

python test.py --template cstlplus --outf ./exp/cstlplus/ --method cstlplus --pretrainedmodelpath ./modelzoo/cst/cstl_plus.pth

GAP_Net

python test.py --template gapnet --outf ./exp/gapnet/ --method gapnet --pretrainedmodelpath ./modelzoo/gapnet/gapnet.pth

ADMM_Net

python test.py --template admmnet --outf ./exp/admmnet/ --method admmnet --pretrainedmodelpath ./modelzoo/admmnet/admmnet.pth

TSA_Net

python test.py --template tsanet --outf ./exp/tsanet/ --method tsanet --pretrainedmodelpath ./modelzoo/tsanet/tsanet.pth

HDNet

python test.py --template hdnet --outf ./exp/hdnet/ --method hdnet --pretrainedmodelpath ./model_zoo/hdnet/hdnet.pth

DGSMP

python test.py --template dgsmp --outf ./exp/dgsmp/ --method dgsmp --pretrainedmodelpath ./model_zoo/dgsmp/dgsmp.pth

BIRNAT

python test.py --template birnat --outf ./exp/birnat/ --method birnat --pretrainedmodelpath ./model_zoo/birnat/birnat.pth

MSTPlusPlus

python test.py --template mstplusplus --outf ./exp/mstplusplus/ --method mstplusplus --pretrainedmodelpath ./modelzoo/mstplusplus/mstplus_plus.pth

λ-Net

python test.py --template lambdanet --outf ./exp/lambdanet/ --method lambdanet --pretrainedmodelpath ./modelzoo/lambdanet/lambdanet.pth

DAUHST_2stg

python test.py --template dauhst --outf ./exp/dauhst2stg/ --method dauhst2stg --pretrainedmodelpath ./modelzoo/dauhst/dauhst2stg.pth

DAUHST_3stg

python test.py --template dauhst --outf ./exp/dauhst3stg/ --method dauhst3stg --pretrainedmodelpath ./modelzoo/dauhst/dauhst3stg.pth

DAUHST_5stg

python test.py --template dauhst --outf ./exp/dauhst5stg/ --method dauhst5stg --pretrainedmodelpath ./modelzoo/dauhst/dauhst5stg.pth

DAUHST_9stg

python test.py --template dauhst --outf ./exp/dauhst9stg/ --method dauhst9stg --pretrainedmodelpath ./modelzoo/dauhst/dauhst9stg.pth ```

  • The reconstrcuted HSI will be output into MST/real/test_code/exp/

The testing scripts in real dataset should be updated <<<

5.3 Visualization

  • Put the reconstruted HSI in MST/visualization/real_results/results and rename it as method.mat, e.g., mstplusplus.mat.

  • Generate the RGB images of the reconstructed HSI

shell cd MST/visualization/ Run show_real.m

 

6. Citation

If this repo helps you, please consider citing our works:

```shell

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

CST

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

DAUHST

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

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

```

Owner

  • Name: BBC
  • Login: bryanbocao
  • Kind: user
  • Location: Dr. Wily's Castle

One of the many Ph.D. candidates.

GitHub Events

Total
Last Year

Dependencies

requirements.txt pypi
  • einops *
  • opencv-python ==4.4.0.46
  • scipy ==1.0.0
  • torch ==1.6.0
  • torchvision ==0.7.0