Universal Image Restoration Pre-training via Masked Degradation
  Classification
Authors: name: JiaKui Hu, name: Zhengjian Yao, name: Lujia Jin, name: Yinghao Chen, name: Yanye Lu
Organization: Peking University

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.