Universal Image Restoration Pre-training via Degradation Classification

TL;DR

This work identifies degradation type information as a valuable, latent prior for pre-training universal image restoration models. It introduces Degradation Classification Pre-Training (DCPT), an encoder–decoder framework where an encoder learns restoration-relevant features and a lightweight degradation classifier provides weak supervision by predicting input degradation; training alternates between degradation classification and image generation, while preserving the encoder's generative capacity. DCPT yields substantial PSNR improvements across all-in-one, single-task, and mixed-degradation scenarios, with gains up to 2.55 dB on 10D all-in-one and over 5 dB in mixed degradation benchmarks, outperforming prior two-stage or large-model baselines. The results support the existence of discriminative priors within restoration models and point to broader opportunities for degradation-aware pre-training to enhance generalization and transfer across tasks and architectures.

Abstract

This paper proposes the Degradation Classification Pre-Training (DCPT), which enables models to learn how to classify the degradation type of input images for universal image restoration pre-training. Unlike the existing self-supervised pre-training methods, DCPT utilizes the degradation type of the input image as an extremely weak supervision, which can be effortlessly obtained, even intrinsic in all image restoration datasets. DCPT comprises two primary stages. Initially, image features are extracted from the encoder. Subsequently, a lightweight decoder, such as ResNet18, is leveraged to classify the degradation type of the input image solely based on the features extracted in the first stage, without utilizing the input image. The encoder is pre-trained with a straightforward yet potent DCPT, which is used to address universal image restoration and achieve outstanding performance. Following DCPT, both convolutional neural networks (CNNs) and transformers demonstrate performance improvements, with gains of up to 2.55 dB in the 10D all-in-one restoration task and 6.53 dB in the mixed degradation scenarios. Moreover, previous self-supervised pretraining methods, such as masked image modeling, discard the decoder after pre-training, while our DCPT utilizes the pre-trained parameters more effectively. This superiority arises from the degradation classifier acquired during DCPT, which facilitates transfer learning between models of identical architecture trained on diverse degradation types. Source code and models are available at https://github.com/MILab-PKU/dcpt.

Figures (7)

• Figure 1: The T-SNE results of randomly initialized PromptIR's feature (left) and all-in-one trained PromptIR's feature (right).
• Figure 2: DCPT follows an encoder-decoder design. The encoder refers to a restoration network, and the decoder is a degradation classifier. DCPT consists of two stages. In each training iteration, (a) degradation classification stage and (b) generation stage occur performed alternately. $L_{cls}$ and $L_{pix}$ are the losses incurred during stage (a) and stage (b). After DCPT, the encoder is fine-tuned for downstream restoration tasks.
• Figure 3: DC-guided training is used for cross-degradation transfer learning. The target task in this figure is denoising.
• Figure 4: Visual comparison on 5D all-in-one image restoration datasets. Zoom in for best view.
• Figure 5: The radar chart of 10D all-in-one image restoration results.