This repository contains the codes for paper End-to-End Low Cost Compressive Spectral Imaging with Spatial-Spectral Self-Attention (ECCV (2020)) by Ziyi Meng*, Jiawei Ma*, Xin Yuan (*Equal contributions). [pdf]
We provide simulation data and real data of our system. You can download them by the following links.
[Simu data (Google Drive)], [Simu data (One Drive)], [Simu data (Baidu Drive pw:kwpn)]
[Real data (Google Drive)], [Real data (One Drive)], [Real data (Baidu Drive pw:j1ii)]
TSA-Net pytorch version is available now. Please find 'TSA_pytorch' file. This code is only for simulation data.
This pytorch model is slightly different from the version in paper, but simpler and more computationally efficient, and achieves similar results.
The pretrained model can be downloaded by the following links. [TSA_pytorch_model (One Drive)], [TSA_pytorch_model (Baidu Drive pw:dba4)]
Moreover, the training data can be downloaded by [Training data (One Drive)], [Training data (Baidu Drive pw:l29y)]
Coded aperture snapshot spectral imaging (CASSI) is an effective tool to capture real-world 3D hyperspectral images. We have proposed a Spatial-Spectral Self-Attention module to jointly model the spatial and spectral correlation in an order-independent manner, which is incorporated in an encoder-decoder network to achieve high quality reconstruction for CASSI.
Fig. 1 (a) Single disperser coded aperture snapshot spectral imaging (SD-CASSI) and our experimental prototype. (b) 25 (out of 28) reconstructed spectral channels. (c) Principle of hardware coding.
Fig. 2 (a) Spatial-Spectral Self-Attention (TSA) for one V feature (head). The spatial correlation involves the modelling for x-axis and y-axis separately and aggregation in an order-independent manner: the input is mapped to Q and K for each dimension: the size of kernel and feature are specified individually. The spectral correlation modelling will flatten samples in one spectral channel (2D plane) as a feature vector. The operation in dashed box denotes the network structure is shared while trained in parallel. (b) TSA-Net Architecture. Each convolution layer adopts a 3 x 3 operator with stride 1 and outputs O-channel cube. The size of pooling and upsampling is P and T.
Fig. 3 A reconstructed hyperspectral video.
Requirements are Python 3 and Tensorflow 1.13.
Download this repository via git
git clone https://github.com/mengziyi64/TSA-Net
or download the zip file manually.
Download the model file
Simulation model(0.97 GB): via Google Drive or One Drive or Baidu Drive (pw:cskg) and put the file into path 'TSA_Net_simulation/Result/Model-Config/TSA-Model/'.
Real data model(1.30 GB): via Google Drive or One Drive or Baidu Drive (pw:iubr) and put the file into path 'TSA_Net_realdata/Result/Model-Config/TSA-Model-Real/'.
Ziyi Meng, Beijing University of Posts and Telecommunications, Email: mengziyi64@163.com
Jiawei Ma, Columbia University, Email: jiawei.m@columbia.edu
Xin Yuan, Bell Labs, Email: xyuan@bell-labs.com