Pré-treinamento Universal para Restauração de Imagens via Classificação de Degradação Mascarada

Universal Image Restoration Pre-training via Masked Degradation Classification

October 15, 2025

Autores: JiaKui Hu, Zhengjian Yao, Lujia Jin, Yinghao Chen, Yanye Lu

cs.AI

Resumo

Este estudo introduz um método de Pré-Treinamento de Classificação de Degradação Mascarada (MaskDCPT), projetado para facilitar a classificação dos tipos de degradação em imagens de entrada, levando a um pré-treinamento abrangente de restauração de imagens. Diferente dos métodos convencionais de pré-treinamento, o MaskDCPT utiliza o tipo de degradação da imagem como uma supervisão extremamente fraca, enquanto simultaneamente aproveita a reconstrução da imagem para melhorar o desempenho e a robustez. O MaskDCPT inclui um codificador e dois decodificadores: o codificador extrai características da imagem de baixa qualidade mascarada. O decodificador de classificação utiliza essas características para identificar o tipo de degradação, enquanto o decodificador de reconstrução visa reconstruir uma imagem de alta qualidade correspondente. Esse design permite que o pré-treinamento se beneficie tanto da modelagem de imagem mascarada quanto do aprendizado contrastivo, resultando em uma representação generalizada adequada para tarefas de restauração. Beneficiando-se do MaskDCPT simples, porém potente, o codificador pré-treinado pode ser usado para abordar a restauração universal de imagens e alcançar desempenho excepcional. A implementação do MaskDCPT melhora significativamente o desempenho tanto para redes neurais convolucionais (CNNs) quanto para Transformers, com um aumento mínimo de 3,77 dB no PSNR na tarefa de restauração all-in-one 5D e uma redução de 34,8% no PIQE em comparação com a linha de base em cenários de degradação do mundo real. Ele também demonstra forte generalização para tipos e níveis de degradação previamente não vistos. Além disso, organizamos e lançamos o conjunto de dados UIR-2.5M, que inclui 2,5 milhões de amostras de restauração pareadas em 19 tipos de degradação e mais de 200 níveis de degradação, incorporando dados sintéticos e do mundo real. O conjunto de dados, código-fonte e modelos estão disponíveis em https://github.com/MILab-PKU/MaskDCPT.

English

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.