LADE-DUN for CASSI

This repo is the implementation of the paper "Latent Diffusion Prior Enhanced Deep Unfolding for Snapshot Spectral Compressive Imaging"

Abstract

Snapshot compressive spectral imaging reconstruction aims to reconstruct three-dimensional spatial-spectral images from a single-shot two-dimensional compressed measurement. Existing state-of-the-art methods are mostly based on deep unfolding structures but have intrinsic performance bottlenecks: i) the ill-posed problem of dealing with heavily degraded measurement, and ii) the regression loss-based reconstruction models being prone to recover images with few details. In this paper, we introduce a generative model, namely the latent diffusion model (LDM), to generate degradation-free prior to enhance the regression-based deep unfolding method by a two-stage training procedure. Furthermore, we propose a Trident Transformer (TT), which extracts correlations among prior knowledge, spatial, and spectral features, to integrate knowledge priors in deep unfolding denoiser, and guide the reconstruction for compensating high-quality spectral signal details. To our knowledge, this is the first approach to integrate physics-driven deep unfolding with generative LDM in the context of CASSI reconstruction. Comparisons on synthetic and real-world datasets illustrate the superiority of our proposed method in both reconstruction quality and computational efficiency.

Architecture

Results Visualization (Real Data)

Usage

Prepare Dataset:

Follow the RDLUF_MixS2, download cave_1024_28 (Baidu Disk, code: fo0q | One Drive), CAVE_512_28 (Baidu Disk, code: ixoe | One Drive), KAIST_CVPR2021 (Baidu Disk, code: 5mmn | One Drive), TSA_simu_data (Baidu Disk, code: efu8 | One Drive), TSA_real_data (Baidu Disk, code: eaqe | One Drive), and then modify the data paths in the option.py.

Pretrained weights and Results

Download pretrained weights and results at (Onedrive).

Environment

python==3.10
torch==2.0.1
scikit-image==0.21.0
scikit-learn==1.5.1
numpy==1.24.4
scipy==1.11.2
pyiqa==0.1.7
matplotlib==3.7.2
Pillow==10.0.0
lpips==0.1.4

Simulation Experiement:

See the readme.md in the ./train_code_syn.

Real Experiment:

See the readme.md in the ./train_code_real.

Acknowledgements

Our code is based on the following codes, thanks for their generous open source:

• https://github.com/ShawnDong98/RDLUF_MixS2
• https://github.com/caiyuanhao1998/MST
• https://github.com/mengziyi64/TSA-Net
• https://github.com/Zj-BinXia/DiffIR

Citation

If this code helps you, please consider citing our works:

@article{wu2023latent,
  title={Latent diffusion prior enhanced deep unfolding for spectral image reconstruction},
  author={Wu, Zongliang and Lu, Ruiying and Fu, Ying and Yuan, Xin},
  journal={arXiv preprint arXiv:2311.14280},
  year={2023}
}