基于掩码降质分类的通用图像修复预训练

本研究介绍了一种掩码退化分类预训练方法（MaskDCPT），旨在促进输入图像的退化类型分类，从而实现全面的图像修复预训练。与传统的预训练方法不同，MaskDCPT使用图像的退化类型作为一种极弱的监督信号，同时利用图像重建来提高性能和鲁棒性。MaskDCPT包含一个编码器和两个解码器：编码器从掩码的低质量输入图像中提取特征。分类解码器利用这些特征来识别退化类型，而重建解码器则旨在重建相应的低质量图像。这种设计使得预训练能够受益于掩码图像建模和对比学习，从而获得适用于修复任务的通用表示。受益于直接而强大的MaskDCPT，预训练的编码器可以用于解决通用的图像修复问题并取得出色的性能。实施MaskDCPT可以显著提高卷积神经网络（CNN）和Transformer的性能，在5D一体化修复任务中，PSNR至少提高了3.77 dB，在真实世界退化场景中，与基线相比，PIQE降低了34.8%。它还表现出对先前未见过的退化类型和级别的强大泛化能力。此外，我们整理并发布了UIR-2.5M数据集，该数据集包含250万个配对的修复样本，涵盖19种退化类型和200多个退化级别，同时包含合成和真实世界数据。数据集、源代码和模型可在https://github.com/MILab-PKU/MaskDCPT获取。