Универсальное предобучение для восстановления изображений через классификацию маскированных искажений

Universal Image Restoration Pre-training via Masked Degradation Classification

October 15, 2025

Авторы: JiaKui Hu, Zhengjian Yao, Lujia Jin, Yinghao Chen, Yanye Lu

cs.AI

Аннотация

В данном исследовании представлен метод предварительного обучения с классификацией замаскированных искажений (MaskDCPT), разработанный для облегчения классификации типов искажений во входных изображениях, что способствует комплексному предварительному обучению восстановлению изображений. В отличие от традиционных методов предварительного обучения, MaskDCPT использует тип искажения изображения в качестве крайне слабого надзора, одновременно используя восстановление изображения для повышения производительности и устойчивости. MaskDCPT включает в себя кодировщик и два декодера: кодировщик извлекает признаки из замаскированного изображения низкого качества. Классификационный декодер использует эти признаки для определения типа искажения, тогда как реконструкционный декодер стремится восстановить соответствующее изображение высокого качества. Такая конструкция позволяет предварительному обучению извлекать выгоду как из моделирования замаскированных изображений, так и из контрастного обучения, что приводит к обобщенному представлению, подходящему для задач восстановления. Благодаря простому, но мощному методу MaskDCPT, предварительно обученный кодировщик может быть использован для решения универсальных задач восстановления изображений и достижения выдающихся результатов. Реализация MaskDCPT значительно улучшает производительность как сверточных нейронных сетей (CNN), так и трансформеров, с минимальным увеличением PSNR на 3.77 дБ в задаче 5D all-in-one восстановления и снижением PIQE на 34.8% по сравнению с базовым уровнем в сценариях реальных искажений. Также наблюдается сильная обобщающая способность к ранее не встречавшимся типам и уровням искажений. Кроме того, мы создали и опубликовали набор данных UIR-2.5M, который включает 2.5 миллиона парных образцов для восстановления по 19 типам искажений и более чем 200 уровням искажений, включая как синтетические, так и реальные данные. Набор данных, исходный код и модели доступны по адресу 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.