Degradation-Aware Residual-Conditioned Optimal Transport for Unified Image Restoration

TL;DR

A Degradation-Aware Residual-Conditioned Optimal Transport (DA-RCOT) approach that models (all-in-one) image restoration as an optimal transport (OT) problem for unpaired and paired settings, introducing the transport residual as a degradation-specific cue for both the transport cost and the transport map.

Abstract

All-in-one image restoration has emerged as a practical and promising low-level vision task for real-world applications. In this context, the key issue lies in how to deal with different types of degraded images simultaneously. In this work, we present a Degradation-Aware Residual-Conditioned Optimal Transport (DA-RCOT) approach that models (all-in-one) image restoration as an optimal transport (OT) problem for unpaired and paired settings, introducing the transport residual as a degradation-specific cue for both the transport cost and the transport map. Specifically, we formalize image restoration with a residual-guided OT objective by exploiting the degradation-specific patterns of the Fourier residual in the transport cost. More crucially, we design the transport map for restoration as a two-pass DA-RCOT map, in which the transport residual is computed in the first pass and then encoded as multi-scale residual embeddings to condition the second-pass restoration. This conditioning process injects intrinsic degradation knowledge (e.g., degradation type and level) and structural information from the multi-scale residual embeddings into the OT map, which thereby can dynamically adjust its behaviors for all-in-one restoration. Extensive experiments across five degradations demonstrate the favorable performance of DA-RCOT as compared to state-of-the-art methods, in terms of distortion measures, perceptual quality, and image structure preservation. Notably, DA-RCOT delivers superior adaptability to real-world scenarios even with multiple degradations and shows distinctive robustness to both degradation levels and the number of degradations.

Figures (10)

• Figure 1: (a) The core idea of DA-RCOT is to model AIR as an OT problem and then condition the OT map with the customized multi-scale residual embeddings, yielding a degradation-aware and structure-preserving DA-RCOT map. (b) A denoising demo under noise level $\sigma=50$. DA-RCOT produces a noise-free image with more faithful textures.
• Figure 2: Overview of the DA-RCOT framework for AIR. DA-RCOT integrates the transport residual into the transport cost, yielding the FROT objective; and more crucially, into the transport map via a two-pass conditioning process. The first pass unconditionally generates an intermediate result along with the residual $\hat{r}_0$. Then the residual is encoded and adapted as multi-scale embeddings $\{\mathbf R_i\}_{i=1,2,3}$ to condition the second-pass restoration.
• Figure 3: Left: Visual examples of the transport residual $r$. Right: Counts of the frequency amplitude of residuals for five types of degradation. For all degradations except for the noise, the residuals are generally sparse in the Fourier domain. The curves are averaged with 40 degraded-clean pairs.
• Figure 4: The t-SNE visual comparison of the degradation embeddings from PromptIR potlapalli2023promptir and DA-RCOT. As compared to the prompt-based embeddings of PromptIR potlapalli2023promptir, the residual embeddings $\mathbf R_1$ of DA-RCOT are clearly separated according to the specific tasks. Particularly, the residual embeddings w.r.t. different levels of noise are clustered together but also exhibit level-specific positional relationships.
• Figure 5: Visual comparison of five-degradation All-in-One results. DA-RCOT restores sharp images with more faithful structural contents.