Transformer based Pluralistic Image Completion with Reduced Information Loss

TL;DR

PUT introduces a patch-based vector-quantized auto-encoder (P-VQVAE) and an Un-Quantized Transformer (UQ-Transformer) to perform pluralistic image inpainting with minimal information loss. By avoiding input downsampling and direct quantization of inputs, PUT preserves high-frequency details and textures, while enabling diverse outputs via a multi-token sampling strategy and optional semantic/structural guidance. The dual-codebook quantization separates masked and unmasked patches, and the multi-scale guided decoder preserves unmasked regions during reconstruction. Experiments on FFHQ, Places2, and ImageNet show PUT achieves superior fidelity and diversity, operates faster than token-by-token autoregressive methods, and allows controllable inpainting through user-provided maps, with promising potential for downstream tasks.

Abstract

Transformer based methods have achieved great success in image inpainting recently. However, we find that these solutions regard each pixel as a token, thus suffering from an information loss issue from two aspects: 1) They downsample the input image into much lower resolutions for efficiency consideration. 2) They quantize $256^3$ RGB values to a small number (such as 512) of quantized color values. The indices of quantized pixels are used as tokens for the inputs and prediction targets of the transformer. To mitigate these issues, we propose a new transformer based framework called "PUT". Specifically, to avoid input downsampling while maintaining computation efficiency, we design a patch-based auto-encoder P-VQVAE. The encoder converts the masked image into non-overlapped patch tokens and the decoder recovers the masked regions from the inpainted tokens while keeping the unmasked regions unchanged. To eliminate the information loss caused by input quantization, an Un-quantized Transformer is applied. It directly takes features from the P-VQVAE encoder as input without any quantization and only regards the quantized tokens as prediction targets. Furthermore, to make the inpainting process more controllable, we introduce semantic and structural conditions as extra guidance. Extensive experiments show that our method greatly outperforms existing transformer based methods on image fidelity and achieves much higher diversity and better fidelity than state-of-the-art pluralistic inpainting methods on complex large-scale datasets (e.g., ImageNet). Codes are available at https://github.com/liuqk3/PUT.

Figures (16)

• Figure 1: Top: Existing transformer based autoregressive methods wan2021highyu2021diverse. The output is produced by ICT wan2021high. Bottom: Our transformer based method. "Tokenize" means getting the indices of quantized pixels or latent vectors, and "De-Tokenize" is the inverse operation.
• Figure 2: Pipeline of PUT for pluralistic image inpainting without condition. The input image is divided into non-overlapped patches which are processed by the encoder independently. The unmasked regions are reconstructed by the decoder, while the masked regions are predicted by the transformer.
• Figure 3: Training procedure of P-VQVAE. The unmasked pixels in the reference image $\mathbf{\hat{x}} \otimes \mathbf{m'}$ are utilized to recover the corresponding pixels in $\mathbf{\hat{x}}^R$, while the latent vectors in codebook $\mathbf{e}'$ and $\mathbf{e}$ are used to recover the pixels in $\mathbf{\hat{x}}^R$ masked by $\mathbf{m}$ and the remaining pixels, respectively.
• Figure 4: The pipeline of controllable image inpainting. Users are free to provide none, one, or both of these two conditions.
• Figure 5: Inpainting results produced by different methods on different datasets. The images produced by PUT are of higher diversity and visual quality.