btkforever/bisci

A Toolbox for Binarized Spectral Compressive Imaging (NeurIPS 2023)

 

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

Qualitative comparisons on the simulation dataset

Qualitative comparisons on the real dataset

Supported Algorithms:

• BiConnect (NeuIPS 2015)
• BNN (NeuIPS 2016)
• Bi-Real (ECCV 2018)
• IRNet (CVPR 2020)
• ReActNet (ECCV 2020)
• BBCU (ICLR 2023)
• BTM (AAAI 2021)
• BiSRNet (NeurIPS 2023)

 

1. Create Environment:

• Python 3 (Recommend to use Anaconda)

• PyTorch >= 1.3

• NVIDIA GPU + CUDA

• Python packages:

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