Skip to content

A Toolbox for Binarized Spectral Compressive Imaging (NeurIPS 2023)

License

Notifications You must be signed in to change notification settings

caiyuanhao1998/BiSCI

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

31 Commits
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

 

arXiv zhihu

PWC

PWC

A Toolbox for Binarized Spectral Compressive Imaging Reconstruction

News

  • 2023.11.27 : Our BiSCI has been added to the Awesome-Model-Quantization collection. 🎉
  • 2023.10.24 : Code, pre-trained models, and results are released. Feel free to check and use them. ⭐
  • 2023.09.22 : Our paper has been accepted by NeurIPS 2023. Code and models will be released. 🚀

Results

  • This repo is a toolbox containing eight BNN-based methods for spectral compressive imaging

  • We provide download links of pre-trained models and results as follows

Simulation results Simulation models Real results Real models
   Google or Baidu       Google or Baidu        Google or Baidu         Google or Baidu    
Quantitative comparisons on the simulation dataset

bnn_compare

quantitative_compare

Qualitative comparisons on the simulation dataset

quantitative_compare

Qualitative comparisons on the real dataset

real_compare

Supported Algorithms:

 

1. Create Environment:

  pip install -r requirements.txt

 

2. 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:

|--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 MST series work, 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.

 

3. Simulation Experiments

3.1 Training

cd simulation/train_code/

# Binnary Connect
python train.py --outf ./exp/biconnect/ --method biconnect

# BNN
python train.py --outf ./exp/bnn/ --method bnn

# Bi-Real Net
python train.py --outf ./exp/bireal/ --method bireal

# IR-Net
python train.py --outf ./exp/irnet/ --method irnet

# ReActNet
python train.py --outf ./exp/reactnet/ --method reactnet

# BBCU
python train.py --outf ./exp/bbcu/ --method bbcu

# BTM
python train.py --outf ./exp/btm/ --method btm

# BiSRNet
python train.py --outf ./exp/bisrnet/ --method bisrnet
  • The training logs, trained models, and reconstructed HSIs will be available in simulation/train_code/exp/

3.2 Testing

Download the pre-trained weights from (Google Drive / Baidu Disk, code: cyh2) and place them to simulation/test_code/model_zoo/

cd simulation/test_code/

# Binnary Connect
python test.py --outf ./exp/biconnect/ --method biconnect --pretrained_model_path ./model_zoo/biconnect.pth

# BNN
python test.py --outf ./exp/bnn/ --method bnn --pretrained_model_path ./model_zoo/bnn.pth

# Bi-Real Net
python test.py --outf ./exp/bireal/ --method bireal --pretrained_model_path ./model_zoo/bireal.pth

# IR-Net
python test.py --outf ./exp/irnet/ --method irnet --pretrained_model_path ./model_zoo/irnet.pth

# ReActNet
python test.py --outf ./exp/reactnet/ --method reactnet --pretrained_model_path ./model_zoo/reactnet.pth

# BBCU
python test.py --outf ./exp/bbcu/ --method bbcu --pretrained_model_path ./model_zoo/bbcu.pth

# BTM
python test.py --outf ./exp/btm/ --method btm --pretrained_model_path ./model_zoo/btm.pth

# BiSRNet
python test.py --outf ./exp/bisrnet/ --method bisrnet --pretrained_model_path ./model_zoo/bisrnet.pth
  • The reconstructed HSIs will be output into simulation/test_code/exp/. Then place the reconstructed results into simulation/test_code/Quality_Metrics/results and run the following commend to calculate the PSNR and SSIM of the reconstructed HSIs.
Run cal_quality_assessment.m
  • Evaluating the Params and FLOPS of models

    We have provided a function my_summary_bnn() in simulation/test_code/utils.py, please use this function to evaluate the parameters and computational complexity of BNNs.

from utils import my_summary_bnn
my_summary_bnn(BiSRNet(), 256, 256, 28, 1)

3.3 Visualization

  • Put the reconstructed HSIs in visualization/simulation_results/results/ and rename it as method.mat, e.g., bisrnet.mat

  • Generate the RGB images of the reconstructed HSIs:

 cd visualization/
 run show_simulation.m 
  • Draw the spectral density lines
cd visualization/
Run show_line.m

 

4. Real Experiments

4.1 Training

cd real/train_code/

# Binnary Connect
python train.py --outf ./exp/biconnect/ --method biconnect

# BNN
python train.py --outf ./exp/bnn/ --method bnn

# Bi-Real Net
python train.py --outf ./exp/bireal/ --method bireal

# IR-Net
python train.py --outf ./exp/irnet/ --method irnet

# ReActNet
python train.py --outf ./exp/reactnet/ --method reactnet

# BBCU
python train.py --outf ./exp/bbcu/ --method bbcu

# BTM
python train.py --outf ./exp/btm/ --method btm

# BiSRNet
python train.py --outf ./exp/bisrnet/ --method bisrnet
  • The training logs, trained models, and reconstructed HSIs will be available in real/train_code/exp/

4.2 Testing

Download the pre-trained weights from (Google Drive / Baidu Disk, code: cyh2) and place them to simulation/test_code/model_zoo/

cd real/test_code/

# Binnary Connect
python test.py --method biconnect --pretrained_model_path ./model_zoo/biconnect.pth

# BNN
python test.py --method bnn --pretrained_model_path ./model_zoo/bnn.pth

# Bi-Real Net
python test.py --method bireal --pretrained_model_path ./model_zoo/bireal.pth

# IR-Net
python test.py --method irnet --pretrained_model_path ./model_zoo/irnet.pth

# ReActNet
python test.py --method reactnet --pretrained_model_path ./model_zoo/reactnet.pth

# BBCU
python test.py --method bbcu --pretrained_model_path ./model_zoo/bbcu.pth

# BTM
python test.py --method btm --pretrained_model_path ./model_zoo/btm.pth

# BiSRNet
python test.py --method bisrnet --pretrained_model_path ./model_zoo/bisrnet.pth
  • The reconstructed HSIs will be output into real/test_code/results/

4.3 Visualization

  • Put the reconstructed HSI in visualization/real_results/results and rename it as method.mat, e.g., bisrnet.mat.

  • Generate the RGB images of the reconstructed HSI

cd visualization/
Run show_real.m

 

5. Citation

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