/TSA-Net

End-to-End Low Cost Compressive Spectral Imaging with Spatial-Spectral Self-Attention

Primary LanguagePython

TSA-Net for CASSI

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 link.
[Simu data (Google Drive)], [Simu data (One Drive)], [Simu data (Baidu Drive pw:aw5u)]
[Real data (Google Drive)], [Real data (One Drive)], [Real data (Baidu Drive pw:)]

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.