Universal Image Restoration Pre-training via Masked Degradation Classification

2025-10-15

AI Summary

This work addresses the challenges of unknown degradation types and poor generalization in generic image restoration pre-training. We propose MaskDCPT, the first method to incorporate degradation-type classification as a weak supervision signal into pre-training, integrating masked image modeling with contrastive learning within an encoder–dual-decoder architecture. The framework jointly optimizes degradation classification loss, image reconstruction loss, and contrastive loss, and is compatible with both CNN- and Transformer-based backbones. To support this research, we release the large-scale UIR-2.5M dataset. Extensive experiments demonstrate that MaskDCPT achieves ≥3.77 dB PSNR improvement on the 5D “all-in-one” restoration task and reduces PIQE by 34.8% on real-world images—substantially outperforming baselines. Moreover, it exhibits strong generalization to unseen degradation types.

Technology Category

• Computer Vision: Diffusion Models for Vision
• Machine Learning: Deep Generative Models & Autoencoders
• Search and Optimization: Learning to Search

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Web Mining and Content Analysis: Large pretrained models with web data
• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking

Abstract

This study introduces a Masked Degradation Classification Pre-Training method (MaskDCPT), designed to facilitate the classification of degradation types in input images, leading to comprehensive image restoration pre-training. Unlike conventional pre-training methods, MaskDCPT uses the degradation type of the image as an extremely weak supervision, while simultaneously leveraging the image reconstruction to enhance performance and robustness. MaskDCPT includes an encoder and two decoders: the encoder extracts features from the masked low-quality input image. The classification decoder uses these features to identify the degradation type, whereas the reconstruction decoder aims to reconstruct a corresponding high-quality image. This design allows the pre-training to benefit from both masked image modeling and contrastive learning, resulting in a generalized representation suited for restoration tasks. Benefit from the straightforward yet potent MaskDCPT, the pre-trained encoder can be used to address universal image restoration and achieve outstanding performance. Implementing MaskDCPT significantly improves performance for both convolution neural networks (CNNs) and Transformers, with a minimum increase in PSNR of 3.77 dB in the 5D all-in-one restoration task and a 34.8% reduction in PIQE compared to baseline in real-world degradation scenarios. It also emergences strong generalization to previously unseen degradation types and levels. In addition, we curate and release the UIR-2.5M dataset, which includes 2.5 million paired restoration samples across 19 degradation types and over 200 degradation levels, incorporating both synthetic and real-world data. The dataset, source code, and models are available at https://github.com/MILab-PKU/MaskDCPT.

Problem

• Classifying degradation types in images for universal restoration tasks
• Enhancing image reconstruction through masked modeling and contrastive learning
• Improving performance across multiple degradation types and levels

Innovation

• Masked degradation classification for image restoration pre-training
• Encoder with dual decoders for classification and reconstruction
• Generalized representation benefiting from masked modeling and contrastive learning