BNN

November 28, 2023 · View on GitHub

 

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. :tada:
  • 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. :rocket:

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 resultsSimulation modelsReal resultsReal 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}
}