[Image Restoration] AOT-GAN "Aggregated Contextual Transformations for High-Resolution Image Inpainting"

contribute

1. Propose aggregated contextual transformations (AOT) for high-resolution image inpainting, which allows capturing informative long-range context and rich patterns of interest for contextual reasoning.
2. A new mask prediction task is designed to train a discriminator for image inpainting, so that the discriminator can distinguish real patches from synthetic patches, thus helping the generator to synthesize fine-grained textures.

model structure

the whole frame

---

AOT block

The generator is first encoded by several standard convolutional layers, then by an AOT block, and finally decoded by a transposed convolution.

AOT blocks adopt a split-transform-merge strategy in three steps:

(1) Splitting: The AOT block splits the convolution kernel of the standard convolution into multiple sub-cores, and each sub-convolution kernel has fewer output channels;

(2) Conversion: Each sub-convolution kernel has a different dilation rate. Larger dilation rates enable sub-kernels to focus on larger regions of the input image, while sub-kernels using smaller dilation rates focus on local patterns with smaller receptive fields.

(3) Aggregation: Context transformations from different receptive fields are finally aggregated by concatenation and standard convolution for feature fusion.

Such a design enables the AOT block to predict each output pixel of the image through different views.

In the following formula, the traditional identical residual connection is improved, and the residual connection is gated. In the aggregation formula, g is a spatially variable gating value. This spatially varying feature aggregation preserves known features outside the missing region while updating features inside the missing region as much as possible.

$$x_3=x_1\times g+x_2\times (1-g)$$

---

Soft Mask-Guided PatchGAN (SM-PatchGAN)

What problem does it solve?

Most deep inpainting models tend to generate an average of all possible solutions based on a reconstruction loss (L1 Loss), which leads to blurry textures.

The inpainting result is expressed as:
$$z=x\odot (1-m)+G(x\odot (1-m),m)\odot The\; result\; of\; m$$
inpainting is the superposition of two parts, the intact area of ​​the original image and the generated hole area. Among them, m is a binary mask (0 means known pixels, 1 means unknown pixels), that is, the missing area is represented as white.

Adversarial loss for the discriminator:

$\begin{array}{c}{L}_{adv}^D={\mathbb{E}}_{z\sim p_z}\left[(D(z)-s(1-m){)}^2\right]+{\mathbb{E}}_{x\sim p_{\text{data }}}\left[(D(x)-1{)}^2\right]\end{array}$

Among them, $\sigma$ represents the composite function of downsampling and Gaussian filtering.

Generator's adversarial loss:

$L_{adv}^G=E_{z\sim p_z}\left[\left(D\right(z)-1{)}^2\odot m\right]$

Design on Discriminator

A soft patch-level mask is designed for the discriminator.

Comparison of different designs:

PatchGAN’s discriminator classifies all patches in the inpainted image as fake, which ignores that the patches outside the missing regions are indeed from real images. While the proposed SM-PatchGAN is able to distinguish the synthetic patches of missing regions from the real patches of the context, which can enhance the discriminator's ability.

No Gaussian filter is used in HM-PatchGAN, thus ignoring that the boundary of the patched image may contain both real and synthetic pixels around the boundary. While the proposed SM-PatchGAN introduces a Gaussian filter to solve this problem.

---

overall optimization

The optimization function includes four: $L_1$loss (reconstruction loss), style loss (style loss), perceptual loss (perceptual loss) and adversarial loss (against loss).

（1）$L_1$loss ensures pixel-level reconstruction accuracy
$$L_{rec}=\Vert x-G(x\odot (1-m),m){\Vert }_1$$
(2) The perceptual loss aims to minimize the $L_1$距离
$$L_{per}=\sum _i\frac{{\Vert {\varphi }_i(x)-{\varphi }_i(z)\Vert }_1}{N_i}$$
Among them, ${\varphi }_i$Activation map from layer i of a pretrained network (such as VGG19), $N_i$It is ${\varphi }_i$The total quantity in .

(3) The style loss is defined as L 1 L_1 between the Gram matrix of the inpainted image and the deep features of the real image$L_1$距离：
$$L_{sty}={\mathbb{E}}_i\left[{\Vert {\varphi }_i(x{)}^T{\varphi }_i(x)-{\varphi }_i(z{)}^T{\varphi }_i\left(z\right)\Vert }_1\right]$$
（4）adversarial loss
$$L_{adv}={\mathbb{E}}_{z\sim p_z}\left[(D(z)-1{)}^2\odot m\right]$$
Overall optimization goal:
$$L=l_{adv}{L}_{adv}^G+l_{rec}{L}_{rec}+l_{per}{L}_{per}+l_{sty}{L}_{sty}$$
Parameter setting: $l_{adv}$= 0.01, $l_{rec}$= 1, $l_{per}$ = 0.1, $l_{sty}$ = 250。

---

implementation details

Gaussian filtering in SM-PatchGAN sets the kernel size of the Gaussian kernel to 70×70. To avoid the color shift problem caused by normalization layers, all normalization layers in the generator network are removed.

Training parameter settings:

In a mini-batch, 8 pictures and corresponding masks are randomly selected. The learning rate of both the generator and the discriminator is $10^{-4}$ , using$b_1=0\text{and}\beta {}_2=0.9$ optimizer. Use VGG19 pre-trained on the ImageNet dataset as a pre-trained network for computing style loss and perceptual loss.

---

experiment

dataset used

Places2、CELEBA-HQ、QMUL-OpenLogo

mask dataset

The mask dataset provided in the paper Image Inpainting for Irregular Holes Using Partial Convolutions is also used in most image restoration tasks.

Model benchmarks for comparison

（1）CA：Context encoders: Feature learning by inpainting. (2016)

（2）PEN-Net：Learning pyramid-context encoder network for high-quality image inpainting. (2019)

（3）PConv：Image inpainting for irregular holes using partial convolutions. (2018)

（4）EdgeConnect：Edgeconnect: Generative image inpainting with adversarial edge learning. (2019)

（5）GatedConv：Free-form image inpainting with gated convolution. (2019)

（6）HiFill ：Contextual residual aggregation for ultra high-resolution image inpainting. (2020)

（7）MNPS ：High-resolution image inpainting using multi-scale neural patch synthesis. (2017)

The above seven models are all classic models in the field of Image Inpainting.

Evaluation Criteria

$L_1$error、PSNR、SSIM、FID

Then there are qualitative experiments, quantitative experiments, and User Study. The results are definitely better than others, so I won’t summarize them. See the paper for details.

Ablation experiment

Verify the effectiveness of the three components in AOT-GAN: gated contextual transformations (gated context conversion), gated residual connections (gated residual connections), SM-PatchGAN discriminator (SM PatchGAN discriminator).

---

in conclusion

limitation

(1) The number of branches and the expansion rate of the AOT block are based on empirical research and settings. When the image size changes, it may be necessary to re-set the parameters and cannot be adaptive.

(2) It is difficult to automatically segment the logo area in practical applications (such as logo removal).