BiSCI: A Toolbox for Binarized Spectral Compressive Imaging

GitHub code arXiv paper NeurIPS 2023 ηŸ₯乎解读

Code and project documentation

Overview

This repository provides pretrained PyTorch checkpoints for binarized spectral compressive imaging reconstruction, including BiSRNet and seven binary neural network baselines supported by the BiSCI toolbox.

The models reconstruct hyperspectral image cubes with 28 spectral bands from coded snapshot measurements. Simulation testing generates the measurements from the provided hyperspectral ground truth and sensing mask; real testing reconstructs the captured measurements directly.

Binarized Spectral Compressive Imaging β€” NeurIPS 2023 introduces the Binarized Spectral-Redistribution Network (BiSRNet). Its BiSR-Conv units redistribute spectral representations before binarization and use a scalable hyperbolic tangent approximation during backpropagation.

Authors: Yuanhao Cai, Yuxin Zheng, Jing Lin, Xin Yuan, Yulun Zhang, and Haoqian Wang.

BiSRNet architecture

Reconstruction examples

Simulation and real reconstruction comparisons from the project:

Simulation reconstruction comparison Real reconstruction comparison

Available checkpoints

There are 12 checkpoints: 8 for simulation and 4 for real experiments. Each checkpoint requires the matching architecture and method setting in the corresponding part of the toolbox.

Method --method Simulation checkpoint Real checkpoint
BiConnect biconnect Simu_models/biconnect.pth Real_models/biconnect.pth
BNN bnn Simu_models/bnn.pth Real_models/bnn.pth
Bi-Real Net bireal Simu_models/bireal.pth β€”
IR-Net irnet Simu_models/irnet.pth β€”
ReActNet reactnet Simu_models/reactnet.pth Real_models/reactnet.pth
BBCU bbcu Simu_models/bbcu.pth β€”
BTM btm Simu_models/btm.pth β€”
BiSRNet bisrnet Simu_models/bisrnet.pth Real_models/bisrnet.pth

The real training code supports all eight architectures. For the four methods without a released real checkpoint here, train a model using the real training code before testing.

Getting started

1. Install the toolbox

Use Python 3.10 and an NVIDIA GPU. The following configuration was tested with PyTorch 2.5.1 and CUDA 12.1 wheels:

git clone https://github.com/caiyuanhao1998/BiSCI.git
cd BiSCI
conda create -n bisci python=3.10 -y
conda activate bisci
python -m pip install torch==2.5.1 torchvision==0.20.1 --index-url https://download.pytorch.org/whl/cu121
python -m pip install -r requirements.txt
python -m pip check

2. Download the model weights

From the BiSCI repository root, download the weights while preserving the checkpoint subdirectories:

# All eight simulation checkpoints
hf download CaiYuanhao/BiSCI --include "Simu_models/*.pth" --local-dir simulation/test_code/model_zoo

# All four real checkpoints
hf download CaiYuanhao/BiSCI --include "Real_models/*.pth" --local-dir real/test_code/model_zoo

To download only the simulation BiSRNet checkpoint:

hf download CaiYuanhao/BiSCI Simu_models/bisrnet.pth --local-dir simulation/test_code/model_zoo

For example, this saves simulation/test_code/model_zoo/Simu_models/bisrnet.pth.

3. Prepare the test data

Follow the dataset preparation instructions to download the test scenes and sensing mask separately. This model repository contains the pretrained weights.

BiSCI/
└── datasets/
    β”œβ”€β”€ TSA_simu_data/
    β”‚   β”œβ”€β”€ mask.mat
    β”‚   └── Truth/
    β”‚       β”œβ”€β”€ scene01.mat
    β”‚       β”œβ”€β”€ ...
    β”‚       └── scene10.mat
    └── TSA_real_data/
        β”œβ”€β”€ mask.mat
        └── Measurements/
            β”œβ”€β”€ scene1.mat
            β”œβ”€β”€ ...
            └── scene5.mat

Simulation test scenes contain img arrays of shape 256 Γ— 256 Γ— 28, and the simulation mask.mat contains a 256 Γ— 256 mask. Real measurement files contain meas_real arrays of shape 660 Γ— 714.

4. Run simulation testing

From the BiSCI repository root:

cd simulation/test_code
python test.py --method bisrnet --pretrained_model_path ./model_zoo/Simu_models/bisrnet.pth --outf ./exp/bisrnet/

The reconstructed cubes are saved in simulation/test_code/exp/bisrnet/Test_result.mat. The file contains pred and truth, both shaped 10 Γ— 256 Γ— 256 Γ— 28.

5. Run real testing

From the BiSCI repository root:

cd real/test_code
python test.py --method bisrnet --pretrained_model_path ./model_zoo/Real_models/bisrnet.pth

The five reconstructions are saved as real/test_code/results/bisrnet/1.mat through 5.mat. Each contains a res array shaped 660 Γ— 660 Γ— 28.

The scripts use GPU 0 by default. Add --gpu_id N to select another GPU. See the GitHub README for other methods, training, metrics, and visualization instructions.

Citation

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

License

The model repository retains its Apache 2.0 license metadata. The toolbox source code is distributed under the MIT license.

Downloads last month

-

Downloads are not tracked for this model. How to track
Inference Providers NEW
This model isn't deployed by any Inference Provider. πŸ™‹ Ask for provider support

Paper for CaiYuanhao/BiSCI