HyperReconNet

PyTorch codes for reproducing the paper: Lizhi Wang, Tao Zhang, Ying Fu, and Hua Huang, HyperReconNet: Joint Coded Aperture Optimization and Image Reconstruction for Compressive Hyperspectral Imaging, TIP, 2019.[Link]

Abstract

Coded aperture snapshot spectral imaging (CASSI) system encodes the 3D hyperspectral image (HSI) within a single 2D compressive image and then reconstructs the underlying HSI by employing an inverse optimization algorithm, which equips with the distinct advantage of snapshot but usually results in low reconstruction accuracy. To improve the accuracy, existing methods attempt to design either alternative coded apertures or advanced reconstruction methods, but cannot connect these two aspects via a unified framework, which limits the accuracy improvement. In this paper, we propose a convolution neural network-based end-to-end method to boost the accuracy by jointly optimizing the coded aperture and the reconstruction method. On the one hand, based on the nature of CASSI forward model, we design a repeated pattern for the coded aperture, whose entities are learned by acting as the network weights. On the other hand, we conduct the reconstruction through simultaneously exploiting intrinsic properties within HSI-the extensive correlations across the spatial and spectral dimensions. By leveraging the power of deep learning, the coded aperture design and the image reconstruction are connected and optimized via a unified framework. Experimental results show that our method outperforms the state-of-the-art methods under both comprehensive quantitative metrics and perceptive quality.

Data

In the paper, two benchmarks are utilized for training and testing. Please check them in Link1(ICVL) and Link2(Harvard). In addition, an extra-experiment following TSA-Net is implemented on CAVE Dataset and KAIST Dataset. To start your work, make HDF5 files of the same length and place them in the correct path. The file structure is as follows:

--data/

--ICVL_train/

--trainset_1.h5
...
--trainset_n.h5
--train_files.txt
--validset_1.h5
...
--validset_n.h5
--valid_files.txt

--ICVL_test/

--test1/
...
--testn/

A few descriptions of datasets can be checked in README. Note that, every image for testing is saved as several 2D images according to different channels.

Environment

Python 3.6.2
CUDA 10.0
Torch 1.7.0
OpenCV 4.5.4
h5py 3.1.0
TensorboardX 2.4
spectral 0.22.4

Usage

Download this repository via git or download the ZIP file manually.

git clone https://github.com/MaxtBIT/HyperReconNet.git

Download the pre-trained models if you need.
Make the datasets and place them in correct paths. Then, adjust the settings in utils.py according to your data.
Run the file main.py to train a model.
Run the files test_for_paper.py and test_for_kaist.py to test models.

Results

1. Reproducing Results on ICVL&Harvard Datasets

The results reproduced on ICVL Dataset and Harvard Dataset. In this stage, the mask is learnable. And the size of patches is 64 * 64 * 31. In addition, only the central areas with 512 * 512 * 31 are compared in testing.

	ICVL		Havard
	Paper	Reproducing	Paper	Reproducing
PSNR	33.63	34.76	31.36	31.39
SSIM	0.990	0.973	0.973	0.900
SAM	0.032	0.040	0.104	0.113

2. Results of Extra-Experiments on CAVE&KAIST Datasets

The results obtained on CAVE Dataset and KAIST Dataset. 30 scenes of CAVE are used for training, and 10 scenes of KAIST are used for testing. The fixed mask is a binary constant randomly generated in model initialization. The optimized mask is a learnable binary variable that can be optimized by the network. Note that, there is only one binary mask utilized in training and testing. Images with a size of 256 * 256 * 28 are used for comparison.

	Fixed Mask	Optimized Mask
PSNR	33.61	36.13
SSIM	0.915	0.950
SAM	0.103	0.075

Citation

@article{HyperReconNet,
  title={HyperReconNet: Joint Coded Aperture Optimization and Image Reconstruction for Compressive Hyperspectral Imaging},
  author={Wang, Lizhi and Zhang, Tao and Fu, Ying and Huang, Hua},
  journal={IEEE Transactions on Image Processing},
  volume={28},
  number={5},
  pages={2257-2270},
  year={2019},
}

robotexplorer/HyperReconNet