Pré-entraînement universel pour la restauration d'images via la classification de dégradations masquées

Universal Image Restoration Pre-training via Masked Degradation Classification

October 15, 2025

papers.authors: JiaKui Hu, Zhengjian Yao, Lujia Jin, Yinghao Chen, Yanye Lu

cs.AI

papers.abstract

Cette étude présente une méthode de pré-entraînement par classification de dégradation masquée (MaskDCPT), conçue pour faciliter la classification des types de dégradation dans les images d'entrée, conduisant à un pré-entraînement complet en restauration d'images. Contrairement aux méthodes de pré-entraînement conventionnelles, MaskDCPT utilise le type de dégradation de l'image comme une supervision extrêmement faible, tout en exploitant simultanément la reconstruction de l'image pour améliorer les performances et la robustesse. MaskDCPT comprend un encodeur et deux décodeurs : l'encodeur extrait les caractéristiques de l'image d'entrée de faible qualité masquée. Le décodeur de classification utilise ces caractéristiques pour identifier le type de dégradation, tandis que le décodeur de reconstruction vise à reconstruire une image de haute qualité correspondante. Cette conception permet au pré-entraînement de bénéficier à la fois de la modélisation d'images masquées et de l'apprentissage contrastif, aboutissant à une représentation généralisée adaptée aux tâches de restauration. Grâce à la simplicité et à la puissance de MaskDCPT, l'encodeur pré-entraîné peut être utilisé pour aborder la restauration universelle d'images et obtenir des performances exceptionnelles. La mise en œuvre de MaskDCPT améliore significativement les performances des réseaux de neurones convolutifs (CNN) et des Transformers, avec une augmentation minimale du PSNR de 3,77 dB dans la tâche de restauration tout-en-un 5D et une réduction de 34,8 % du PIQE par rapport à la référence dans des scénarios de dégradation du monde réel. Elle montre également une forte généralisation à des types et niveaux de dégradation précédemment inconnus. En outre, nous avons constitué et publié le jeu de données UIR-2.5M, qui comprend 2,5 millions d'échantillons de restauration appariés couvrant 19 types de dégradation et plus de 200 niveaux de dégradation, intégrant à la fois des données synthétiques et réelles. Le jeu de données, le code source et les modèles sont disponibles à l'adresse 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.