Source-linked AI summary
JPEG Artifact Correction using Denoising Diffusion Restoration Models
Bahjat Kawar, Jiaming Song, Stefano Ermon, Michael Elad
TL;DR
Existing diffusion-based inverse-problem solvers are mostly limited to linear operators, while JPEG artifact correction requires handling a nonlinear degradation. The paper generalizes DDRM’s pseudo-inverse for nonlinear measurement operators and uses a pretrained unconditional diffusion model for restoration. Across JPEG quality factors, the method performs on par with a state-of-the-art JPEG-specific baseline while generalizing better to lower quality factors and image dequantization without retraining.
Problem
Existing inverse-problem approaches generally target linear operators, leaving nonlinear problems such as JPEG artifact correction insufficiently covered.
Method
The method extends DDRM by generalizing its pseudo-inverse concept to other measurement operators, using JPEG encoding and decoding within an unconditional diffusion restoration procedure.
Results
Across quality factors, the method compares favorably with a JPEG-restoration GAN baseline on PSNR, SSIM, and LPIPS [41], with greater improvement at low QF outside that baseline’s training range.
Takeaways & Limitations
The approach generalizes beyond JPEG to image dequantization and does so without retraining or problem-specific hyperparameter tuning.
Takeaways & Limitations
The method relies on quantization matrices being stored within JPEG files and therefore known at inference time.
Abstract
from arXiv · showhide
Diffusion models can be used as learned priors for solving various inverse problems. However, most existing approaches are restricted to linear inverse problems, limiting their applicability to more general cases. In this paper, we build upon Denoising Diffusion Restoration Models (DDRM) and propose a method for solving some non-linear inverse problems. We leverage the pseudo-inverse operator used in DDRM and generalize this concept for other measurement operators, which allows us to use pre-trained unconditional diffusion models for applications such as JPEG artifact correction. We empirically demonstrate the effectiveness of our approach across various quality factors, attaining performance levels that are on par with state-of-the-art methods trained specifically for the JPEG restoration task.
1 Introduction
Existing diffusion-based inverse-problem methods generally target linear operators, leaving nonlinear tasks such as JPEG artifact correction insufficiently addressed. The paper extends DDRM with a generalized pseudo-inverse and applies it to JPEG restoration using an unconditional diffusion model.
- Motivation: Existing problem-agnostic inverse-problem methods generally handle linear operators, whereas JPEG artifact correction is a nonlinear inverse problem.Training separate models for multiple degradation types can be expensive, motivating inference-time adaptation without retraining.
- Method: The method generalizes DDRM’s pseudo-inverse from linear operators to broader measurement operators, including JPEG encoding.For JPEG, decoding serves as the corresponding pseudo-inverse, yielding an update that replaces the linear operator and pseudo-inverse with JPEG encoding and decoding.
- Results: Across quality factors, the method compares favorably with a JPEG-specific GAN baseline on PSNR, SSIM, and LPIPS [41].It improves especially at low quality factors not covered by the baseline’s training, supporting the generalization advantage of an unconditional diffusion model.
2 Background
The paper reviews diffusion models, DDRM, linear inverse problems, and JPEG compression as foundations for diffusion-based restoration. These foundations motivate using pretrained unconditional diffusion models rather than problem-specific conditional models.
- Diffusion Models: Diffusion models generate samples through a Markov chain from x_T to x_0, retaining x_0 as the final output.Their denoising parameterization can be trained with a denoising autoencoder objective.
- Linear Inverse Problems: A linear inverse problem seeks to recover x from measurements y = Hx + z, where H is known and z is additive white Gaussian noise with known variance.The formulation distinguishes the unknown signal, measurements, degradation matrix, and measurement noise.
- Related Work: Unconditional diffusion-based inverse-problem solvers avoid retraining for each degradation model, unlike conditional models trained on paired signals and measurements.This problem-agnostic formulation can save computational resources and improve applicability across inverse problems.
- Denoising Diffusion Restoration Models (DDRM): DDRM solves linear inverse problems by transforming signals and measurements into a shared spectral space using the singular value decomposition of H.It denoises dimensions with nonzero singular values and imputes dimensions with zero singular values while accounting for measurement noise.
- JPEG: JPEG compresses images through color conversion, optional chroma subsampling, 8 × 8 blocking, DCT, quantization, and lossless entropy coding.Information is lost during chroma subsampling and quantization, while decoding reverses the overall process.
3 JPEG Artifact Correction with DDRM
The method extends DDRM beyond linear operators by replacing the linear pseudo-inverse with JPEG decoding, enabling diffusion-based JPEG restoration.
- 3 JPEG Artifact Correction with DDRM: Generalizing DDRM’s pseudo-inverse properties to nonlinear operators enables JPEG artifact correction with JPEG decoding as the operator’s pseudo-inverse.JPEG encoding preserves the re-encoded measurement, while decoding after encoding produces an image close to the original.
- 3 JPEG Artifact Correction with DDRM: The JPEG restoration update replaces DDRM’s linear operator and pseudo-inverse with JPEG encoding and decoding, respectively.The resulting process injects measurement-derived information and combines it with the current input and denoiser output at each iteration.
- 3 JPEG Artifact Correction with DDRM: The method is practical because JPEG quantization matrices are stored in the JPEG files and are therefore known at inference time.This allows the restoration update to use the relevant compression information without requiring unavailable metadata.
4 Experimental Results
Experiments on ImageNet-1K evaluate the method across JPEG quality factors and show favorable restoration quality against decoding and a specialized baseline, alongside generalization to image dequantization.
- JPEG artifact correction: Across JPEG quality factors, the method significantly improves over simple decoding and performs favorably or comparably to QGAC on PSNR, SSIM, and LPIPS.QGAC is trained for QF 10–100, whereas the proposed method generalizes across all QF without JPEG-specific training.
- JPEG artifact correction: The method generalizes well to low JPEG quality factors where the specifically trained QGAC baseline generalizes poorly.The experiments use JPEG quality factors from 1 to 100, with lower values representing heavier compression.
- Image dequantization: The method also generalizes to image dequantization, recovering high-quality reconstructions for images quantized below the standard 24 bits per color.This extends the demonstrated scope beyond JPEG artifact correction to another nonlinear inverse problem.
5 Conclusion
The paper extends DDRM to nonlinear inverse problems with a generalized pseudo-inverse and demonstrates JPEG restoration that remains competitive while generalizing to lower quality factors and image dequantization.
- 5 Conclusion: The proposed diffusion method extends DDRM beyond linear inverse problems by generalizing the pseudo-inverse concept.It is evaluated on ImageNet-1K for JPEG restoration and further demonstrated on image dequantization.
- 5 Conclusion: On ImageNet-1K, the method performs on par with a state-of-the-art baseline in most cases and generalizes to lower quality factors that specifically trained baselines do not.The approach also generalizes beyond JPEG without retraining or problem-specific hyperparameter tuning.
A Additional Visual Results
Visual examples compare ground-truth, JPEG-compressed, and restored images across quality factors, with zoomed regions highlighting artifact removal.
- A Additional Visual Results: Across different quality factors, restored images are presented alongside ground-truth and JPEG-compressed images to show correction of compression artifacts.Zoomed-in regions in the lower-right corners highlight specific artifact removals.