Stimulating Diffusion Model for Image Denoising via Adaptive Embedding and Ensembling

---

Abstract: Image denoising is a fundamental problem in computational photography, where achieving high perception with low distortion is highly demanding. Current methods either struggle with perceptual quality or suffer from significant distortion. Recently, the emerging diffusion model has achieved state-of-the-art performance in various tasks and demonstrates great potential for image denoising. However, stimulating diffusion models for image denoising is not straightforward and requires solving several critical problems. For one thing, the input inconsistency hinders the connection between diffusion models and image denoising. For another, the content inconsistency between the generated image and the desired denoised image introduces distortion. To tackle these problems, we present a novel strategy called the Diffusion Model for Image Denoising (DMID) by understanding and rethinking the diffusion model from a denoising perspective. Our DMID strategy includes an adaptive embedding method that embeds the noisy image into a pre-trained unconditional diffusion model and an adaptive ensembling method that reduces distortion in the denoised image. Our DMID strategy achieves state-of-the-art performance on both distortion-based and perception-based metrics, for both Gaussian and real-world image denoising.

---

🚀 News

• 2023.5 The first version of the manuscript waw submitted and the code was uploaded. ⭐

• 2024.5.10 Our paper has been accepted by TPAMI! 🎉

• 2024.6.4 The code and the tools are all supplemented and released! 🎊

• 2024.11.2 The detailed steps and tips for testing with other datasets are provided! 💗

• 2024.11.8 The citation imformation is updated. 🌹

• 2024.11.16 Fix some minor problems, making the code more beginner-friendly. 🌞

• 2024.11.23 Release a comparison-friendly version: DMID-4500. 🌈

If you find this repo useful, please give it a star ⭐ and consider citing our paper in your research. Thank you.

⏳ Todo lists

• We will supplement the code about noise transformation within a month (before 6.12).
• We may release our other methods ...
• We may release more results of commonly used datasets ...
• We may release our reproduced methods, for example R2R, ZS-N2N ...

Pipeline of DMID

Quick Start

• Download the pre-trained unconditional diffusion model(from OpenAI) and place it in ./pre-trained/.

• To quick start, just run:

python main_for_gaussian.py

Evaluation

• All the visual results are available.
• Download testsets (CBSD68, Kodak24, McMaster, Urban100, ImageNet), and place the testsets in './data/', eg './data/CBSD68'.
• Download the testsets after noise transformation (CC, PolyU, FMDD), and replace the folder named '.pre-trained' with the downloaded testsets.
• Download the pre-trained unconditional diffusion model(from OpenAI) and place it in ./pre-trained/.

• To quickly reproduce the reported results, run

sh evaluate.sh

Tools

• 🔨 All the details of setting details of different tables can be found in the paper.
• 🔨 To calculate timestep N for a given noise level, we provied two version code, you can find them in https://github.com/Li-Tong-621/DMID/utils

python utils_cal_N.py

python utils_cal_N_2.py

• 🔨 To perform our improved noise transformation method by yourself, or denoised any give noisy image, please firstly perform noise transformation and then denoise the intermediate image.

python main_for_real_NT.py
python main_for_real.py

• The detailed tips are as follows:

• • 1: Note that in python main_for_real_NT.py, we need to set some parameters, such as the noise level parameter sigma related to the noise level of the images (line 430: sigma=data_dct['sigma']), and the parameter eps for determining the stopping condition (line 445: eps=data_dct['eps']).
• • 2: Running python main_for_real_NT.py successfully will generate a .pt file, after which we can run the main_for_real.py file. When running main_for_real.py, we need to find an appropriate step like line 155: diffusion_times = find_N(sigma=2 * noises[str(index)]['noise_level'] / 255).
• • 3: In general, choosing appropriate parameters can yield very good results. For easier operation or comparison, feel free to make some changes, for example, use a fixed stopping condition, such as running 4500 or 5000 iterations fixed for each image.

• 🔨 We provide a new code for real-world image denoising (main_for_real.py), because there are some errors, which i didn't find, in original code for real-world image denoising (main_for_real_o.py).

• 🔨 What if the denoised images you have are not natural images?

• • 1: You can first try to denoise directly with the current model. We also used non-natural image datasets in our paper, for example, the FMDD dataset, which is actually a dataset of fluorescent cell images. We treated the FMDD dataset as single-channel grayscale natural images, and you can refer to the code for specific operations.
• • 2: You can also train a diffusion model with your own dataset. We used a model trained by OpenAI on ImageNet, and you can use OpenAI's code to train a generative model for polarized images. Some potential issues you might encounter can be found in the issues #9 and #10.

Results

Classical Gaussion denoising

Robust Gaussion denoising

Real-world image denoising

Compared with other diffusion-based methods

Additional Results

Some researchers have inquired about the results of the SIDD dataset. We briefly degrad the noise transformation in the embedding method to a fixed 4500 iterations and remove the ensembling method. The results are referred to DMID-4500:

Method	SIDD validation	SIDD benchmark	CC	FMDD
	PSNR / SSIM / LPIPS	PSNR / SSIM / LPIPS	PSNR / SSIM / LPIPS	PSNR / SSIM / LPIPS
DMID-4500	33.41 / 0.9130 / 0.343	34.21 / 0.8482 / -	35.70 / 0.9811 / 0.139	31.68 / 0.8298 / 0.421

Citation

@inproceedings{DMID,
	title={Stimulating the Diffusion Model for Image Denoising via Adaptive Embedding and Ensembling},
	author={Li, Tong and Feng, Hansen and Wang, Lizhi and Xiong, Zhiwei and Huang, Hua},
	booktitle={IEEE Transactions on Pattern Analysis and Machine Intelligence (TPAMI)},
	pages={8240-8257},
	year={2024},
}