Fully Context-Aware Image Inpainting with a Learned Semantic Pyramid

In recent years, many deep learning models based on the architectures of conditional GAN have been proposed for deterministic image inpainting [12, 13, 1, 14, 15]. These models also greatly benefit from the adversarial training scheme as well as further designs. Pathak et al. [16] first exploited the auto-encoder framework to perform global context encoding and image decoding. Iizuka et al. [17] proposed to use multiple discriminators based on stacked dilated convolutional layers to enhance the consistency of global and local visual context. Other architectures are also explored in recent work such as the multi-scale learning modules [5, 18]. Unlike the existing work, our model specifically leverages the multi-scale information as the semantic prior pyramid, which is learned in a multi-level knowledge distillation module.

More recent studies take advantage of the attention module in capturing long-range contextual relations [13, 19, 20]. For instance, Deepfill [13] follows a coarse-to-fine generation framework with patch-wise aggregation to enhance the restoration of high-frequency information. Besides, some methods propose various network architectures to eliminate the undesired influence from invalid pixels in context feature extraction [21, 4]. Although these approaches perform generally better than rule-based approaches at generating realistic images, they still struggle with structural distortion in complex scenes. This is because they cannot deal well with the ambiguous context of the missing area due to the limitation of the deterministic training paradigm.

To cope with the multi-modal data distribution between the corrupted input images and the expected output images, different probabilistic models have been proposed, which can model the probability of the potential content in the missing area and produce diverse results. PICNet [22] exploits the framework of VAE in two different network branches to learn the conditional distribution within the missing areas. UCTGAN [23] also uses a similar architecture with the manifold projection module. Also, there are GAN-based methods [24, 25] and auto-regressive methods [26] that explicitly consider the diversity of generated images. Unlike all of the above methods, our model is particularly designed to harness the multi-modal distribution by leveraging a plug-and-play variational inference module in the learning process of semantic priors, instead of directly using it in the image generator.

Some state-of-the-art approaches also use explicit visual priors as the learning supervision to improve the quality of image inpainting. Nazeri et al. [6] and Xiong et al. [12] presented two-stage frameworks that first learn to generate the images of object edges under external supervisions, and then use them to guide texture restoration. Some methods also use color-smooth images [27, 7, 28] or semantic maps [29] to provide more dense, structural, and meaningful supervisions. However, these methods either suffer from the accumulated errors through the two processing stages, or require heavy human annotations, which is difficult to scale to real-world applications. In contrast, first, the semantic priors in our model are jointly learned with the image generator, which greatly reduce error accumulation; Second, the priors are learned from the pyramid of feature maps of a pretext model, which are more accessible. The concept of knowledge distillation was first proposed in the work from Hinton et al. [30] to transfer pre-learned knowledge from a large teacher network to a smaller student network. Although Suin et al. [31] also used a joint-training autoencoder as distillation target, it lacks pre-learned semantic knowledge which are important for structural guidance. Here, we adapt the knowledge distillation framework to high-level semantic priors and low-level texture priors simultaneously from the pretext model, both of which are important for context-aware image inpainting.

As shown in Fig. 2(a), the semantic prior learner follows the U-Net architecture. It distills multi-scale visual priors from a particular pretext model given a full image $\bm{I}_{\text{full}}\!\in\!\mathbb{R}^{3\times h\times w}$, a damaged image $\bm{I}_{\text{mask}}\!\in\!\mathbb{R}^{3\times h\times w}$, and a corresponding binary mask $\bm{M}\!\in\!\mathbb{R}^{1\times h\times w}$ where $1$ represents the locations for invalid pixels and $0$ represents those for valid pixels. Considering that the global scene understanding usually involves multiple objects, in this work, we borrow a pretext multi-label classification model, denoted by $H$, to provide feature maps as the supervisions of learning the semantic priors. Unless otherwise specified, the pretext model is trained on the OpenImage dataset [33] with the asymmetric loss (ASL) [32]. Note that we do not fine-tune the model on any inpainting datasets used in this paper. We hope the pretext model can provide more discriminative representations on specific objects, allowing the common knowledge to be transferred to the unsupervised image inpainting task. To generate the supervisions of knowledge distillation, specifically, we up-sample the full image $\bm{I}_{\text{full}}\!\in\!\mathbb{R}^{3\times h\times w}$ to obtain $\bm{I}_{\text{full}}^{\prime}\!\in\!\mathbb{R}^{3\times 2h\times 2w}$, and then feed $\bm{I}_{\text{full}}^{\prime}$ to $H$ to obtain $L$ feature maps in total at multiple scales:

where $\bm{S}_{\text{target}}^{l}$ has a spatial size of $\frac{h}{2^{l-1}}\times\frac{w}{2^{l-1}}$. $L$ denotes the number of scales of the used semantic priors. In all experiments, we set $L=3$ and the spatial size in each scale is $256\times 256$, $128\times 128$, and $64\times 64$, respectively.

At the beginning of the semantic prior learner, $L$ down-sampling layers are used, and a residual blockis applied after the last down-sampling layer to extract features from the masked image at different scales. To align with the spatial resolution of $\bm{S}_{\text{target}}^{l}$, we also up-sample $\bm{I}_{\text{mask}}$ into $\bm{I}^{\prime}_{\text{mask}}\!\in\!\mathbb{R}^{3\times 2h\times 2w}$. The corresponding mask $\bm{M}$ is also up-sampled to $\bm{M}^{\prime}$ in the same manner, and concatenated with $\bm{I}^{\prime}_{\text{mask}}$ to form the input of the semantic prior learner:

where $\bm{S}^{l}_{\text{enc}}$ is the image representation learned from visible image contents that has the same spatial size as the semantic supervision $\bm{S}_{\text{target}}^{l}$. We here adopt the common practice about $\bm{M}$ in previous literature, keeping the region of $\bm{M}$ known during training and testing and focusing more on the recovery of missing content. Next, we exploit multiple residual blocks and up-sampling layers to generate a pyramid of semantic priors with multi-scale skip connections:

where $U$ indicates the pixel-shuffle layer for up-sampling and $\{\bm{S}^{l}_{\text{prior}}\}_{l=1}^{L}$ indicates the prior pyramid. For each $\bm{S}^{l}_{\text{prior}}$, it has the spatial sizes of $\frac{h}{2^{l-1}}\times\frac{w}{2^{l-1}}$, and is followed by a distillation head $f^{l}$ of a $1\times 1$ convolutional layer. Finally, we apply a set of $\ell_{1}$ distillation losses to the deterministic semantic prior learner:

where $\odot$ denotes the Hadamard product, $\bm{M}^{l}$ is the resized mask with spatial size equal to $\bm{S}_{\text{prior}}^{l}$, and $\alpha$ is for the additional constraints on masked areas. The above objective functions are jointly used with final image reconstruction loss. In this way, the knowledge at different levels of the pretexts model can flow into the prior pyramid to best facilitate image inpainting. Notably, the distillation heads $\{f^{l}\}_{l=1}^{L}$ and the pretext model $H$ are only used in the training phase and are all removed from SPN during testing.

To resolve the contextual ambiguity, it is natural to consider various possibilities of the missing content in image inpainting. Therefore, in addition to learning a prior pyramid, we further improve our previous work at IJCAI’21 [8] by introducing a variational inference module in the semantic learner. This module is optional and handles the multi-modal distribution of possible missing content in the high-level feature space conditioned on the remaining part of the image. By injecting this module in the semantic learner, we can easily extend SPN to probabilistic image inpainting, that is, generating diverse restoration results given an individual damaged image. More specifically, we reformulate the semantic prior learner based on a Gaussian stochastic neural network. In practice, as shown in Fig. 2(a), we simply embed a variational inference module between the high-level image representation $\bm{S}^{l}_{\text{enc}}$ and the corresponding learned priors $\bm{S}^{L}_{\text{prior}}$. The details of this module are shown in Fig. 3(a). It consists of an encoder $E_{v}$ that estimates the parameters of the posterior distribution of high-level semantics $q_{s}$ and a generator $G_{v}$ that produces diverse prior pyramids with the re-parameterization trick:

The multiple-level distillation losses are then used in the same way as the learning deterministic prior pyramid, which is shown in Eq. (4). We also optimize the KL divergence to constrain the distance between the learned posterior distribution and a prior distribution $p_{z}\sim\mathcal{N}(\bm{0},\bm{1})$:

where $q_{s}$ represents the learned posterior distribution.

By optimizing the KL divergence within the semantic prior learner, our model is very different from existing approaches for stochastic image inpainting in the following perspective: It incorporates the distribution of potential semantics as a part of the learned priors, in the sense that different knowledge is distilled from the pretext model in a probabilistic way.

We use the reconstruction and adversarial losses to train SPN for deterministic inpainting. The reconstruction loss is in an $\ell_{1}$ form and constrains the pixel-level distance between the ground truth images and the inpainting results, with more attention on the missing content:

The adversarial loss, as adopted in existing approaches [4, 6], is used to encourage realistic inpainting results:

where $D$ represents the discriminator and $\mathcal{L}^{D}_{\text{adv}}$ encourages the discriminator to distinguish real and generated images. By combining $\mathcal{L}_{\text{adv}}$ with the aforementioned loss terms, the full objective function can be written as

where the form of $\mathcal{L}_{\text{prior}}$ is shown in Eq. (4). The grid search range for $\lambda_{1}$ is $[0.1,1,10]$ and the range for $\lambda_{2}$ is same as $\lambda_{1}$. We finally set $\lambda_{1}=10$, $\lambda_{2}=1$, and $\delta=4$. Similarly, we set $\alpha=3$ in Eq. (4).

Except for the loss funtions in Eq. (12), we further use the feature matching loss, the perceptual loss and the perceptual diversity loss [24] to train the entire model for probabilistic inpainting.

The feature matching loss measures the distance between feature maps extracted from $N$ layers of the discriminator given real and generated images. We use $\phi^{i}$ to denote the mapping function from the input image to the feature map at the $i$-th layer in the discriminator. For the perceptual loss, we use a fixed VGG19 model pre-trained on ImageNet to extract multi-scale visual representations and reduce their distance in the semantic space. We use $\varphi^{i}$ to denote the feature map from the $i$-th layer in the VGG19 model and $K$ is the total number of layers we use. Combined with the feature matching loss that has similar forms, the two loss terms can be summarized as

As for the perceptual diversity loss, it encourages SPN to produce semantically diverse results in the missing areas. It is also conducted on the feature maps from the pre-trained VGG19 model:

where $\widehat{\bm{I}}_{1}$ and $\widehat{\bm{I}}_{2}$ denote restored images generated from the same damaged image but with different noises, $\bm{M}_{i}$ is the resized mask which has the same spatial size as $\varphi^{i}$, and $\epsilon$ is the perturbation term to avoid the outlier. The full objective function for probabilistic image inpainting can be written as

The grid search range for both $\lambda_{3}$ and $\lambda_{4}$ is $[0.1,1,10]$ and the range for $\lambda_{5}$ is $[0.01,0.05,0.1]$. Finally, we have $\lambda_{3}=10$, $\lambda_{4}=1$, and $\lambda_{5}=0.05$. We refer to SPADE [9] to set other parameters including $K=5$, $N=3$, and $\epsilon=1e^{-5}$.

We validate the effectiveness of SPN on four datasets in total. Under the deterministic image inpainting setup, we evaluate SPN on three datasets. For the probabilistic inpainting setup, we use the CelebA-HQ [34] and Paris StreetView dataset [35].Furthermore, we use an external dataset to provide irregular masks. In both training and testing phases, for CelebA, we follow the common practice [4, 5] to perform center crop on images and then resize them to $256\times 256$. For other datasets, raw images are resized to $256\times 256$ directly.

• 1.

  Places2 [36]. It is one of the most challenging datasets containing over $1.8$ million images from $365$ different scenes. We use the standard training set with $1.8$ million images for training and perform evaluations on the first $12{,}000$ images in the validation set.

• 2.

  CelebA [37]. It contains about $200{,}000$ diverse celebrity facial images. we use the original training set that includes over $162{,}000$ images for training and use the first $12{,}000$ images in the test set for testing.

• 3.

  Paris StreetView [35]. It collects images from street views of Paris and mainly contains different highly structured facades. We use $14{,}900$ images for training and $100$ images for testing. We use this dataset under both deterministic and probabilistic setups.

• 4.

  CelebA-HQ [34]. It is a human facial dataset used by previous probabilistic inpainting methods [22, 24]. It contains a higher quality of texture detail than CelebA. Following the existing work [24], we use about $27{,}000$ images for training and $2{,}800$ images for testing.

• 5.

  Irregular masks [38]. This dataset is used to provide irregular masks in our experiments. As a common practice [4], we perform random flipping at training time. A total number of $12{,}000$ irregular masks are grouped by six intervals of the mask ratio: $0\%$-$10\%$, $10\%$-$20\%$, $\ldots$, $50\%$-$60\%$.

We extend our previous work by providing stronger baseline models. For the deterministic setup, we compare SPN with RFR [5], RN [4], MFE [7], EC [6], WF [39], and CTSDG [28]. Particularly, we denote our previous work as SPL [8], short for the Semantic Prior Learner. Among these methods, RN, RFR, and WF introduce different network architectures to improve both encoding and decoding stages. On the other hand, EC, CTSDG, and MFE try to use edge maps or smooth images as additional structural supervisions to help image completion. However, they still suffer from distorted structures due to the lack of global semantic understanding. For probabilistic inpainting, we compare SPN with PIC [22] and DSI [26].

In our experiments, we extract feature maps at three scales from the multi-label classification model to form the distillation targets. The classification model was pre-trained on the OpenImage dataset [33] with the asymmetric loss (ASL) [32], and is not fine-tuned on any datasets used in this work. Besides, we also use the patch-based discriminator for adversarial training as in the previous work [4, 7]. Our model is trained by Adam solver with $\beta_{1}=0.0$ and $\beta_{2}=0.9$. The initial learning rate is set to $1e^{-4}$ for all experiments, and we decay the learning rate to $1e^{-5}$ in the last quarter of the training process. SPN is trained on two V100 GPUs, and it takes about one day to train our model on Paris StreetView dataset. For more details, please refer to our code and pre-trained models at https://github.com/WendongZh/SPN.

We have the following observations from Table 1. First, the proposed SPN achieves the best results in most evaluation metrics on all three datasets. It not only outperforms its previous version SPL, but also obtains significant improvements (up to $1$dB for PSNR) compared with other methods. Second, the improvements from our model are consistent across the three datasets. Particularly, SPN obtains the most significant improvements over other competitors such as EC and MFE on the Place2 dataset, which demonstrates that SPN can handle more difficult image scenarios. Third, for FID, our model greatly improves the results of its previous work SPL, achieving the best performance among all compared models. This improvement indicates the effectiveness of the newly proposed multi-scale prior learner and fully context-aware image generator in better capturing global context semantics.

Fig. 4 gives the showcases of the restored images on three datasets. For the Places2 dataset which contains complex scenes and various objects, our model successfully generates more complete global and local structures for missing regions. In other columns, we can see that the results of RFR still suffer from distorted structures, while SPN can produce more complete structures as well as clearer details. Particularly, for the last column, we perform object detection on the restored images to show whether the compared models generate images by effectively understanding the semantics of the scene. We can see that, compared with other methods, the generated images from SPN lead to correct object detection results with higher confidence scores.

We also provide more samples with highlighted details in Fig 5. We can observe that results from SPN contain more clear and sharp local details, which further demonstrate the effectiveness of the proposed semantic pyramid.

The evaluation process in probabilistic inpainting is similar to that in deterministic inpainting. Since probabilistic models can generate diverse plausible results given a single input image, we select images with the best PSNR performance from $5$ different samples to calculate the quantitative results. From Table 2, we can observe that SPN significantly outperforms other probabilistic inpainting approaches in all evaluation metrics on different datasets. Although SPN underperforms our previous deterministic approach SPL [8] on most metrics, it still achieves better results on FID with mask ratio $40\%-60\%$. This result can be explained from two aspects. First, SPN pays extra capacities to learn the multi-model distribution of the generated image, which results in the performance degeneration especially in pixel-level metrics. Second, the learned prior pyramid provides consistent multi-scale semantic guidance, which contributes to the restoration of large missing areas.

Fig. 6 shows four restored images by each compared model given the same input images on the CelebA-HQ dataset. We can observe that SPN generates more reasonable content than previous approaches. Specifically, images generated by PIC lack diversity among different samples, and results from DSI suffer from distorted structures such as the ears in the upper two rows and the eyes in the bottom two rows. Instead, human faces restored by SPN not only contain plausible details but also has diverse and reasonable local structures. These results show that the stochastic prior pyramid in our model can successfully model the multi-modal distribution of the potential content in the missing area and provide explicit guidance for both global and local restoration.

We independently remove different modules of SPN and provide results in Fig. 7. First, we remove all semantic supervisions of knowledge distillation from the prior learner, termed as “w/o ASL” (the green bars). By comparing with the blue bars, we observe a significant decline in the performance of SPN. It demonstrates that the pretext knowledge plays an important role in the context-aware image inpainting, and has been effectively transferred to inpainting task.

Besides, we evaluate the effect of multi-scale features in the prior pyramid by only integrating $\bm{S}^{L}_{\text{prior}}$ with the smallest spatial size ($64\times 64$) into the image generator, termed as “w/o Prior Pyramid” (the yellow bars). In Fig. 7, we can observe that replacing the multi-scale priors with single-scale priors leads to a large performance decline. In Table 3, we further provide a detailed ablation on the usage of semantic priors at different scales. We can observe that as we learn the model with more semantic priors, the FID score is significantly improved. These results show that using multi-scale semantic pyramid indeed help inpainting model generate more perceptually realistic results.

Finally, we show the effectiveness of the SPADE ResBlocks by simply concatenating image features and the corresponding prior features in the image generator. For example, we have $\bm{F}^{L}_{\text{dec}}=\texttt{concat}(\bm{F}^{L}_{\text{enc}},\bm{S}^{L}_{\text{prior}})$ in Eq. (9). The results are shown by the grey bars and marked as “w/o SPADE”. We may conclude that the use of SPADE ResBlocks effectively facilitates the joint modeling of both vision features and semantic representations. We also conduct ablations on using different numbers of SPADE ResBlocks in Table 4. We can observe that using too less blocks or too many blocks all degenerate the model performance.

Since other types of structural priors have also been used by existing inpainting methods, we here explore the alternatives of the semantic supervisions , and see whether they can also facilitate context understanding. Specifically, as shown in Fig. 8, we first replace the multi-scale features from the ASL model, whose results are represented by the blue bars, with those from a detection network YoloV5 [41] pre-trained on the MS COCO dataset [42] (the green bars). Similarly, we also take supervisions from a semantic segmentation model HRNet [43] pre-trained on the ADE20K dataset [44] (the grey bars), as well as the edge maps used in the EC method (the yellow bars). Notably, these results cannot be strictly compared with each other due to different pre-training datasets and model architectures. However, we can still obtain some interesting observations. First, the model that uses edge maps as supervisions achieves the worst performance compared with other models using different pre-trained feature maps. Since the edge maps only contain low-level sparse structures instead of high-level semantic information, this result demonstrates that the semantic modeling process can provide more effective guidance for context understanding. Second, object-level tasks such as detection and multi-label classification may result in more discriminative representations for scene understanding.

To further evaluate the effectiveness of learning semantic priors from pre-learned networks, we use t-SNE algorithm to visualize the semantic concepts of the generated results. We replace the originally used supervisions with two different features: edge maps and feature maps extracted by a randomly initialized classification model, named “SPN w / edge” and “SPN w / rand” In Fig 9, we can observe that results from SPN have similar cluster patterns with the real data. On the contrary, results from other ablations may suffer from semantic ambiguities. These results show that with the help of semantic modeling, our approach can generate images with more clear semantic concepts.

Since image inpainting can also be utilized as a self-supervised representation learning method [16], we thus conduct an extra experiment to answer whether the pre-trained SPN model can also provide meaningful semantic representations for image inpainting. Specifically, we use SPL [8], the preliminary approach of SPN, as the baseline model, and replace the original distillation target with $\bm{S}_{\text{prior}}^{L}$ from the pre-trained SPN model on the Places2 dataset. The experiment is conducted on the Paris StreetView dataset. From Table 5, we can see that using pre-trained $\bm{S}_{\text{prior}}^{L}$ as the prior distillation target achieves comparable results with using ASL [32]. It even performs better in the FID score. Besides, as $\bm{S}_{\text{prior}}^{L}$ is only pre-trained on the Places2 dataset, it also shows that the learned representation have potential generalization abilities.

Fig. 10 compares the structural information generated by the proposed semantic prior learner and the edge maps from the EC model. Specifically, we directly perform the K-Means clustering algorithm on the output feature maps and use different colors to represent different clusters. We set $K=8$ for all visualization results and release the clustering function at our Github page. From this figure, we observe that the edge maps from the EC model cannot effectively restore reasonable structures for complex scenes, resulting in distorted structures in the final results. Instead, the prior learner can capture more complete context understanding and provide consistent structural information at different spatial scales. We notice that higher-level semantic priors are more discriminative and are helpful for contextual reasoning, while lower-level priors enhance the understanding of local structure and texture details.

Besides, in Fig. 11, we visualize randomly sampled priors in our variational inference module. Given the same input image, we visualize two priors sampled from different latent variables with the largest spatial size ($256\times 256$). We can observe that the sampled priors contain different structural layout, resulting in diverse details in final output images. These results demonstrate the effectiveness of our variational inference module in handling probabilistic inpainting.

In Fig. 12, we further visualize the features in the image generator to show how the feature refinement works by integrating the learned semantic priors. We here use $\bm{F}^{L}_{\text{enc}}$ and $\bm{F}^{L}_{\text{dec}}$ at the smallest spatial scale ($64\times 64$), which are extracted before and after using the $8$ SPADE ResBlocks. We also visualize $\bm{S}_{\text{prior}}^{L}$ for ease of comparison. We can see that $\bm{F}^{L}_{\text{enc}}$ is shown to focus more on local texture consistency, while $\bm{F}^{L}_{\text{dec}}$ is effectively refined by the learned semantic priors $\bm{S}_{\text{prior}}^{L}$ to capture both global context and local texture.