Residual-Conditioned Optimal Transport: Towards Structure-Preserving Unpaired and Paired Image Restoration

Xiaole Tang,Xin Hu,Xiang Gu,Jian Sun

2024-05-11

Abstract:Deep learning-based image restoration methods generally struggle with faithfully preserving the structures of the original image. In this work, we propose a novel Residual-Conditioned Optimal Transport (RCOT) approach, which models image restoration as an optimal transport (OT) problem for both unpaired and paired settings, introducing the transport residual as a unique degradation-specific cue for both the transport cost and the transport map. Specifically, we first formalize a Fourier residual-guided OT objective by incorporating the degradation-specific information of the residual into the transport cost. We further design the transport map as a two-pass RCOT map that comprises a base model and a refinement process, in which the transport residual is computed by the base model in the first pass and then encoded as a degradation-specific embedding to condition the second-pass restoration. By duality, the RCOT problem is transformed into a minimax optimization problem, which can be solved by adversarially training neural networks. Extensive experiments on multiple restoration tasks show that RCOT achieves competitive performance in terms of both distortion measures and perceptual quality, restoring images with more faithful structures as compared with state-of-the-art methods.

Computer Vision and Pattern Recognition

What problem does this paper attempt to address?

The problem that this paper attempts to solve is how to faithfully preserve the structural details of an image while minimizing the distortion metric in image restoration tasks. Traditional image restoration methods usually rely on designing optimization problems to utilize the prior knowledge of clear images, but these methods can often only capture the "average" value of the high - quality image distribution given the degraded image distribution, resulting in the restored image being too smooth and the structural details being damaged. Although generative methods can produce visually appealing results, they usually use the degraded image as a conditional input without including specific degradation information, which may lead to distortion and inaccurate structural details still existing in the output image.

For this reason, the paper proposes a new Residual - Conditioned Optimal Transport (RCOT) method. By introducing a degradation - related transport residual as a unique clue for transport cost and transport mapping, it aims to solve the above challenges. Specifically, the RCOT method first incorporates degradation - specific information into the transport cost through the Fourier - residual - guided optimal transport objective, and then designs a two - stage RCOT mapping. In the first stage, the base model calculates the transport residual and encodes it into a degradation - specific embedding to regulate the restoration process in the second stage. Through this mechanism, RCOT can dynamically inject degradation - specific knowledge from the residual embedding into the restoration operator, enhancing its ability to preserve the image structure.

The main contributions of the paper include:
1. Modeling image restoration as an optimal transport problem and introducing the Fourier - residual - guided optimal transport objective, which allows the incorporation of degradation - specific knowledge into the transport cost. Further, the minimax dual form of the optimal transport model is derived.
2. Proposing a two - stage RCOT method. By conditioning the transport mapping on the residual embedding, it dynamically injects degradation - specific information from the residual embedding into the restoration operator, enhancing its ability to preserve the image structure.
3. Extensive experiments on multiple tasks (such as image denoising, super - resolution, rain removal, and fog removal) show that this method performs well in both distortion metric and perceptual quality, especially in restoring the structural details of the image, and has obvious advantages compared with existing methods.