Universal Image Restoration Pre-training via Degradation Classification

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.

• The paper introduces Degradation Classification Pre-Training (DCPT), a framework utilizing degradation type as a weak supervisory signal to enhance universal image restoration models.
• DCPT achieves significant performance gains, with up to 2.55 dB in all-in-one tasks and 6.53 dB in mixed degradation scenarios, outperforming previous pre-training methods.
• The framework improves model transferability and generalization across diverse and unseen degradation types and is compatible with both CNN and transformer architectures.

The paper "Universal Image Restoration Pre-training via Degradation Classification" introduces the Degradation Classification Pre-Training (DCPT) framework aimed at enhancing image restoration models' performance by pre-training them to classify degradation types. DCPT is distinctive in its utilization of the degradation type as a weakly supervised signal, readily available in image restoration datasets, to bolster the pre-training process traditionally dominated by self-supervised methods.

Key Contributions:

1. Degradation Classification Framework:
   • DCPT is structured explicitly to enhance models' ability to discern various degradation types. This involves training an encoder to extract features, followed by a lightweight decoder, such as ResNet18, tasked with classifying the degradation type based solely on these features.
   • Uniquely, this methodology uses the degradation classification task to instill prior knowledge in models, enhancing their capacity for universal image restoration tasks.
2. Performance Metrics:
   • DCPT achieves significant performance improvements: an average gain of up to 2.55 dB in a 10D all-in-one restoration task and 6.53 dB in mixed degradation scenarios.
   • A comparative advantage over prior methods is demonstrated, including masked image modeling techniques that often discard the decoder post-pre-training, whereas DCPT retains and effectively utilizes pre-trained parameters.
3. Transferability and Generalization:
   • The paper highlights the superior transfer learning capabilities facilitated by the degradation classifier obtained during DCPT, which enhances model performance across diverse degradation types.
   • It underscores the innate degradation identification ability within models, thus enabling efficient model generalization and previously unseen degradation types recognition.
4. Model Compatibility:
   • The pre-training process effectively benefits both Convolutional Neural Networks (CNNs) and transformers, establishing DCPT’s broad applicability in different architecture paradigms.
5. Ablation Studies and Comparisons:
   • The paper provides detailed ablation studies illustrating the impact of multi-level feature extraction and the architectural choices in the decoder on the pre-training outcomes, alongside comparisons with other degradation-aware training strategies like masked modeling.

In conclusion, the work presents DCPT as an efficient pre-training strategy that capitalizes on minimal supervision from degradation types to construct a robust, generalized image restoration framework. The approach shows promise not only in improving universal restoration capabilities across tasks but also in significantly enhancing models' utility in practical, mixed-degradation conditions. The paper suggests that exploring a model's latent degradation discrimination capability offers a novel direction for universal image restoration methodologies.