Modeling Mask Uncertainty in Hyperspectral Image Reconstruction

This repository contains the source code, pre-trained models and benchmark testing data for the ECCV2022 Oral paper Modeling mask uncertainty in hyperspectral image reconstruction by Jiamian wang, Yulun Zhang, Xin Yuan, Ziyi Meng, and Zhiqiang Tao.

Updates

More pre-trained models of compared methods will come soon! 🚀🚀

Introduction

Recently, hyperspectral imaging (HSI) has attracted increasing research attention, especially for the ones based on a coded aperture snapshot spectral imaging (CASSI) system. Existing deep HSI reconstruction models are generally trained on paired data to retrieve original signals upon 2D compressed measurements given by a particular optical hardware mask in CASSI, during which the mask largely impacts the reconstruction performance and could work as a “model hyperparameter” governing on data augmentations. This mask-specific training style will lead to a hardware miscalibration issue, which sets up barriers to deploying deep HSI models among different hardware and noisy environments. To address this challenge, we introduce mask uncertainty for HSI with a complete variational Bayesian learning treatment and explicitly model it through a mask decomposition inspired by real hardware. Specifically, we propose a novel Graph-based Self-Tuning (GST) network to reason uncertainties adapting to varying spatial structures of masks among dif- ferent hardware. Moreover, we develop a bilevel optimization framework to balance HSI reconstruction and uncertainty estimation, accounting for the hyperparameter property of masks. Extensive experimental results validate the effectiveness (over 33/30 dB) of the proposed method under two miscalibration scenarios and demonstrate a highly competitive performance compared with the state-of-the-art well-calibrated methods.

Figure 1. Illustration of modeling mask uncertainty with the proposed Graph-based Self-Tuning (GST) network

Performance

Table 1. PSNR(dB)/SSIM by different methods on 10 simulation scenes under the many-to-many hardware miscalibration. All the methods are trained with a mask set and tested by random unseen masks. TSA-Net, GSM, and SRN are obtained with a mask ensemble strategy. We report mean/std among 100 testing trials.

Requirements

• Python 3.7.10
• Pytorch 1.9.1
• Numpy 1.21.2
• Scipy 1.7.1

Train

To train from stratch, run the train.py once determining the configurations

python train.py

To resume from a checkpoint upon the last_train epoch, lauch the program at the checkpoint directory. Keep or modify the stop_criteria accordingly, for example, run

python train.py --model_type ST --data_type 24chl --mode many_to_many --last_train 654  --stop_criteria 500

Test

Ten 256x256x28 and ten 256x256x24 benchmark testing data are provided. Testing trials can be determined by specify trial_num

For the 256x256x28 data type (employed in the main paper), pre-trained models of traditional/miscalibration scenarios are provided. Please specify last_train as 623 for mode one_to_one and one_to_many. Specify last_train as 661 for many_to_many. For example, run

python test.py --mode many_to_many --trial_num 100 --last_train 661 --test_data_type 28chl --model_type GST --test_path ./Data/testing/28chl/ --inter_channels 28 --spatial_scale 4 --noise_act softplus

python test.py --mode one_to_many --trial_num 100 --last_train 623 --test_data_type 28chl --model_type GST --test_path ./Data/testing/28chl/ --inter_channels 28 --spatial_scale 4 --noise_act softplus

python test.py --mode one_to_one  --last_train 623 --test_data_type 28chl --model_type GST --test_path ./Data/testing/28chl/ --inter_channels 28 --spatial_scale 4 --noise_act softplus

For the 256x256x24 data type, pre-trained model of many_to_many miscalibration scenario (primary concern) is provided. For example, run

python test.py --test_path ./Data/testing/24chl/test.mat --mode many_to_many --trial_num 100 --last_train 654 --test_data_type 24chl --model_type ST --noise_act softplus 

Structure of directories

directory	description
Data	Ten 256x256x28 testing data and ten 256x256x24 testing data
test	testing script
utils	utility functions
network	GST network and simplified version (ST)
ssim_torch	function for computing SSIM
model	pre-trained models for both 28-channel and 24-channel HSI data
train	training script

Citation

If you find the code helpful in your resarch, please kindly cite the following papers.

@article{wang2021calibrated,
  title={Calibrated Hyperspectral Image Reconstruction via Graph-based Self-Tuning Network},
  author={Wang, Jiamian and Zhang, Yulun and Yuan, Xin and Meng, Ziyi and Tao, Zhiqiang},
  journal={arXiv preprint arXiv:2112.15362},
  year={2021}
}

@article{wang2021new,
  title={A new backbone for hyperspectral image reconstruction},
  author={Wang, Jiamian and Zhang, Yulun and Yuan, Xin and Fu, Yun and Tao, Zhiqiang},
  journal={arXiv preprint arXiv:2108.07739},
  year={2021}
}

Contact

If you have any questions, please contact Jiamian Wang ([email protected]).

Acknowledgements

We refer to the TSA-Net when we develop this code. Great thanks to them!