Skip to content

Latest commit

 

History

History
81 lines (54 loc) · 4.31 KB

README.md

File metadata and controls

81 lines (54 loc) · 4.31 KB

LADE-DUN for CASSI

(ECCV 2024 - Best Paper Award Candidate & Oral)

This repo is the official implementation of the paper titled "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:

Citation

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

@inproceedings{wu2024latent,
  title={Latent Diffusion Prior Enhanced Deep Unfolding for Snapshot Spectral Compressive Imaging},
  author={Wu, Zongliang and Lu, Ruiying and Fu, Ying and Yuan, Xin},
  booktitle={European Conference on Computer Vision},
  pages={164--181},
  year={2024},
  organization={Springer}
}