(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".
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.
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.
Download pretrained weights and results at (Onedrive).
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
See the readme.md in the ./train_code_syn.
See the readme.md in the ./train_code_real.
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
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}
}