Skip to content

PyTorch codes for reproducing TIP 2019 paper "HyperReconNet: Joint Coded Aperture Optimization and Image Reconstruction for Compressive Hyperspectral Imaging".

Notifications You must be signed in to change notification settings

MaxtBIT/HyperReconNet

Repository files navigation

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

  1. Download this repository via git or download the ZIP file manually.
git clone https://github.com/MaxtBIT/HyperReconNet.git
  1. Download the pre-trained models if you need.
  2. Make the datasets and place them in correct paths. Then, adjust the settings in utils.py according to your data.
  3. Run the file main.py to train a model.
  4. 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},
}

About

PyTorch codes for reproducing TIP 2019 paper "HyperReconNet: Joint Coded Aperture Optimization and Image Reconstruction for Compressive Hyperspectral Imaging".

Resources

Stars

Watchers

Forks

Releases

No releases published

Packages

No packages published

Languages