A Modular Conditional Diffusion Framework for Image Reconstruction

Magauiya Zhussip,Iaroslav Koshelev,Stamatis Lefkimmiatis

2024-11-09

Abstract:Diffusion Probabilistic Models (DPMs) have been recently utilized to deal with various blind image restoration (IR) tasks, where they have demonstrated outstanding performance in terms of perceptual quality. However, the task-specific nature of existing solutions and the excessive computational costs related to their training, make such models impractical and challenging to use for different IR tasks than those that were initially trained for. This hinders their wider adoption, especially by those who lack access to powerful computational resources and vast amount of training data. In this work we aim to address the above issues and enable the successful adoption of DPMs in practical IR-related applications. Towards this goal, we propose a modular diffusion probabilistic IR framework (DP-IR), which allows us to combine the performance benefits of existing pre-trained state-of-the-art IR networks and generative DPMs, while it requires only the additional training of a relatively small module (0.7M params) related to the particular IR task of interest. Moreover, the architecture of the proposed framework allows for a sampling strategy that leads to at least four times reduction of neural function evaluations without suffering any performance loss, while it can also be combined with existing acceleration techniques such as DDIM. We evaluate our model on four benchmarks for the tasks of burst JDD-SR, dynamic scene deblurring, and super-resolution. Our method outperforms existing approaches in terms of perceptual quality while it retains a competitive performance with respect to fidelity metrics.

Computer Vision and Pattern Recognition

What problem does this paper attempt to address?

The paper attempts to address two major limitations of existing Diffusion Probabilistic Models (DPMs) in Image Restoration (IR) tasks:

1. **High computational cost**: Existing DPMs require a large number of Neural Function Evaluations (NFEs) during training and inference, making them computationally expensive when dealing with high-resolution images, especially in the absence of powerful computational resources.
2. **Task specificity**: Existing DPMs are usually trained for specific IR tasks, and if they need to be applied to different IR tasks, the entire model needs to be retrained, which is not only time-consuming but also requires a large amount of data.

To address these issues, the authors propose a modular conditional diffusion framework (DP-IR) with the following features:

- **Modular design**: By combining pre-trained state-of-the-art IR networks and generative DPMs, only a relatively small module (about 0.7M parameters) needs to be additionally trained to adapt to different IR tasks.
- **Accelerated sampling strategy**: By introducing a new sampling strategy, the number of Neural Function Evaluations can be reduced by at least four times without loss of performance. Additionally, this strategy can be combined with existing acceleration techniques (such as DDIM) to further improve efficiency.

The authors evaluated the proposed model on four benchmark datasets, including burst joint demosaicking, denoising, and super-resolution (JDD-SR), dynamic scene deblurring, and 4x single image super-resolution (SISR). Experimental results show that the method outperforms existing methods in terms of perceptual quality and is also competitive in fidelity metrics.