Swin-CasUNet: Cascaded U-Net with Swin Transformer for Masked Face Restoration

Abstract

Masked face restoration is one of the most valuable challenges in the computer vision community. With the in-depth study of u-shaped architectures, also known as U-Net, great progress has been achieved in the development of masked face restoration during the past few years. However, previous restoration methods fail to fully model the long-range dependency due to the locality of convolution layers of the U-Net. To address this problem, we propose a shifted windows Transformer (Swin Transformer) based cascaded U-Net framework called Swin-CasUNet, which incorporates the long-range dependency merit of Transformer into the cascaded U-Net architecture to effectively enhance the functionality and generalization of U-shaped architecture. Specifically, we design a two-stage cascaded U-Net architecture to implement the coarse-to-fine restoration of the masked face. Swin Transformers is adopted to extract global self-attention contexts for the feature map produced by the encoder part of the U-Net. An improved face structure loss is proposed to supervise structure learning. To evaluate the robustness of our masked face restoration model, we collect 3800 pairs of full face images and corresponding masked face images from the real-world and web. Experiments on the datasets demonstrate that our proposed method can generate high quality restoration results. In order to quantitatively compare with previous face restoration methods, we modify the input of our system by manually adding regular and irregular white masks on CelebA face datasets, and then retrain our network. Experiments show that our Swin-CasUNet outperforms previous methods on benchmark datasets.