Test-Time Degradation Adaptation for Open-Set Image Restoration

TL;DR

The paper tackles open-set image restoration by addressing unknown degradations that diverge from training data. It introduces TAO, a three-component framework that combines a degradation-agnostic diffusion backbone (PDM), a test-time degradation adapter (TDA), and adapter-guided image restoration (AIR) to enable label-free, test-time adaptation. By aligning generative outputs to the degraded test domain and applying stage-wise guidance during diffusion denoising, TAO achieves competitive performance across dehazing, low-light enhancement, and denoising without task-specific retraining. The results demonstrate that open-set IR can be effectively tackled with test-time adaptation in diffusion-based pipelines, offering practical robustness for real-world restoration tasks.

Abstract

In contrast to close-set scenarios that restore images from a predefined set of degradations, open-set image restoration aims to handle the unknown degradations that were unforeseen during the pretraining phase, which is less-touched as far as we know. This work study this challenging problem and reveal its essence as unidentified distribution shifts between the test and training data. Recently, test-time adaptation has emerged as a fundamental method to address this inherent disparities. Inspired by it, we propose a test-time degradation adaptation framework for open-set image restoration, which consists of three components, \textit{i.e.}, i) a pre-trained and degradation-agnostic diffusion model for generating clean images, ii) a test-time degradation adapter adapts the unknown degradations based on the input image during the testing phase, and iii) the adapter-guided image restoration guides the model through the adapter to produce the corresponding clean image. Through experiments on multiple degradations, we show that our method achieves comparable even better performance than those task-specific methods. The code is available at https://github.com/XLearning-SCU/2024-ICML-TAO.

Figures (7)

• Figure 1: The differences of image restoration (IR) tasks. To be specific, (a) Classic IR works in a close-set scenario where the training and test degradations are the same and known, and customizes a specialized model for each one. (b) All-in-one IR also works in a close-set scenario where the training and test degradations are the same but unknown, and addresses them through a unified model. (c) Zero-shot IR focuses on recovery from single degraded image, which is free from the training degradations, but often requires priors about the test degradations in advance. In contrast, (d) OIR works in an open-set scenario where the test degradations are unknown and different from the pretraining ones. This is analogous to the challenge in natural language processing, which applies the pre-trained large language model to the various downstream tasks not predefined during the pretraining phase.
• Figure 2: Overview of the proposed method, which exploits (i) a PDM as the generic pre-trained model for OIR. After each denoising step, it first performs (ii) TDA for adapting to the unknown and unseen degradations posed by open-set scenarios, and then conducts (iii) AIR for optimizing the guided image towards the restored clean image. Note: the snowflake icon indicates the image or model is fixed, and the flame icon indicates the image or model will be updated through the gradients.
• Figure 3: Qualitative results on image dehazing, from which one could observe that existing methods excessively dehazing resulting in darkening and/or artifacting of the images. In contrast, our method obtains clearer results which are closer to the natural ground truths.
• Figure 4: Qualitative results on low-light image enhancement, from which one could see that our results are not as smooth as MBLLEN nor as dark as ZDEC. Although there are slight color biases from ground truths, our method achieves a rational lighting of the dark images.
• Figure 5: Qualitative results on image denoising, from which one could observe that existing methods excessively denoising resulting in smoothing and/or artifacting of the images. In contrast, our method obtains clearer and sharper results which are closer to ground truths.