TSA-Net for CASSI

March 19, 2024 · View on GitHub

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)]

New

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)]

Overviewer

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.

TSA-Net Architecture

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.

Usage

Download the TSA-Net repository and model file

  1. Requirements are Python 3 and Tensorflow 1.13.
  2. Download this repository via git
git clone https://github.com/mengziyi64/TSA-Net

or download the zip file manually.

  1. 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/'.

Testing

  1. Testing on simulation data
    Run TSA_Net_simulation/test.py to reconstruct 10 synthetic datasets. The results will be saved in 'TSA_Net_simulation/Result/Testing-Result/' in the MatFile format.
  2. Testing on real data
    Run TSA_Net_realdata/test.py to reconstruct 5 real datasets. The results will be saved in 'TSA_Net_realdata/Result/Testing-Result/' in the MatFile format.

Training

  1. Training simulation model
    1. Put hyperspectral datasets (Ground truth) into corrsponding path, i.e., 'TSA_Net_simulation/Data/Training_truth/' for training data and 'TSA_Net_simulation/Data/Valid_truth/' for validation data. For our setting, the training data and validation datashould be scaled to 0-65535 and 0-1, respectively, with a size of 256×256×28.
    2. Adjust training parameter by modify TSA_Net_simulation/Model/Config.yaml.
    3. Run TSA_Net_simulation/train.py.
  2. Training real data model
    1. Put hyperspectral datasets (Ground truth) into corrsponding path, i.e., 'TSA_Net_realdata/Data/Training_truth/' for training data and 'TSA_Net_realdata/Data/Valid_truth/' for validation data. For our setting, the training data and validation datashould be scaled to 0-1 with a size of 660×660×28.
    2. Adjust training parameter by modify TSA_Net_realdata/Model/Config.yaml.
    3. Run TSA_Net_realdata/train.py.

Contact

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