DriftRec: Adapting diffusion models to blind JPEG restoration

TL;DR

This work proposes an elegant modification of the forward stochastic differential equation of diffusion models to adapt them to solve blind JPEG restoration at high compression levels and names it DriftRec, and shows that it can escape the tendency of other methods to generate blurry images, and recovers the distribution of clean images significantly more faithfully.

Abstract

In this work, we utilize the high-fidelity generation abilities of diffusion models to solve blind JPEG restoration at high compression levels. We propose an elegant modification of the forward stochastic differential equation of diffusion models to adapt them to this restoration task and name our method DriftRec. Comparing DriftRec against an $L_2$ regression baseline with the same network architecture and state-of-the-art techniques for JPEG restoration, we show that our approach can escape the tendency of other methods to generate blurry images, and recovers the distribution of clean images significantly more faithfully. For this, only a dataset of clean/corrupted image pairs and no knowledge about the corruption operation is required, enabling wider applicability to other restoration tasks. In contrast to other conditional and unconditional diffusion models, we utilize the idea that the distributions of clean and corrupted images are much closer to each other than each is to the usual Gaussian prior of the reverse process in diffusion models. Our approach therefore requires only low levels of added noise and needs comparatively few sampling steps even without further optimizations. We show that DriftRec naturally generalizes to realistic and difficult scenarios such as unaligned double JPEG compression and blind restoration of JPEGs found online, without having encountered such examples during training.

Figures (9)

• Figure 1: Reproduced from richter2022journal: The forward and reverse process for the ouve sde illustrated with a single scalar variable. The mean (thick black line) of the forward process exponentially decays from a clean sample $\mathbf x_0$ (blue) towards the corrupted sample $\mathbf y$ (green), and the standard deviation (shaded gray region) increases exponentially. The reverse process then starts from a slightly mismatched distribution $\tilde{p}_T$ which is centered around $\mathbf y$ rather than $\mathbf x_T$ and moves to an estimate $\hat{\mathbf{x}_0}$. Five realizations of both processes are shown as thin black lines.
• Figure 2: Mean and variance curves of the forward processes described by the sde \ref{['eq:ouve-sde', 'eq:tsdve-sde', 'eq:cosve-sde']}. The mean curves exhibit different shapes between $\mathbf{x}_0$ and $\mathbf{y}$ but start and end at the same points. The variance curves are all drawn but indistinguishable.
• Figure 3: Example reconstructions on the D2KF2K test set for JPEG quality factor 10, comparing FBCNN jiang2021towards, I2SB liu2023i2sb, the regression baseline, and our proposed method DriftRec. Reconstructions by FBCNN and the baseline exhibit a blurry, painting-like look, whereas DriftRec generates high-frequency detail better matching that of the ground-truth image. I2SB leaves significant blocky artifacts in this blind setting. Best viewed zoomed in.
• Figure 4: Example reconstructions on CelebA-HQ256 for JPEG quality factor 10, comparing FBCNN jiang2021towards, ddrm kawar2022jpeg and the regression baseline against our method. Best viewed zoomed in.
• Figure 5: Example images for unaligned double JPEG compression with quality factors (top) QF$_1$=10, QF$_2$=30 and (bottom) QF$_1$=30, QF$_2$=10. Only DriftRec achieves reliable blind removal of artifacts and exhibits best perceptual quality. "FBCNN@10" indicates FBCNN jiang2021towards with a forced (and thus non-blind) quality factor of 10.