Incremental Transformer Structure Enhanced Image Inpainting with Masking Positional Encoding

In plain words

Image inpainting fills in the missing part of a picture so that it blends with the rest. Many fill-in methods paint plausible texture but get the large structures wrong. They also struggle with big or oddly shaped holes and with large images.

This paper's fix is to split the job in two. First, a transformer draws a rough sketch of the missing area: its outlines and straight lines, on a small grayscale grid. A transformer can look across the whole picture, so it is good at getting the overall structure right. Second, a separate network paints colour and texture, guided by that sketch. The sketch is added to the painter gently, starting from nothing, so the existing painter still works at the start of training. The sketch can also be enlarged to fit bigger images.

The paper adds one more hint for the painter: for each missing pixel, how far it is from known pixels and in which direction. That helps large holes, where a painter otherwise loses track of where it is and repeats patterns.

TL;DR

• Two-stage inpainting. A transformer restores edge and line sketches at 256 px. A FFC-based painter (LaMa architecture) takes the sketch through zero-initialised residual additions (ZeroRA), so a pretrained painter can be extended without retraining from scratch.
• A masking positional encoding (MPE) gives each masked pixel a distance and a direction to the nearest known pixel §3.3.
• Author-reported: at 256 px with mixed masks, PSNR/SSIM/FID/LPIPS beat LaMa on Indoor and Places2 tab. 1. At 512 and 1024 px, the authors report gains over LaMa tab. 5.
• Use it if you need structure-aware inpainting for large or irregular holes and can afford a second network. Expect a sketch-stage failure mode on distant urban views app. G.

Problem

The authors list four limits of existing inpainters §1:

1. CNNs have small receptive fields, so they miss textures and structures that span large holes.
2. Holistic structure (long edges, lines) is hard to recover without global context.
3. Training GANs at large image sizes is costly, and quality drops at high resolution.
4. Masked regions carry no explicit position information, so large irregular holes repeat meaningless patterns.

Auxiliary-structure methods (edges, segmentation, gradients) help with structure but usually need several stages or several models, and adding new inputs means training from scratch §1, §2. The paper's aim is to add structure to an existing pretrained inpainter without retraining it §1, §3.3.

The method assumes the structure stage works at 256 px, which the authors name as the main limit app. G.

Method

Overview §3 fig. 2. The pipeline has these parts:

1. Transformer Structure Restoration (TSR) §3.1. Input is the masked image, Canny edges and lines, and the mask, all at 256 px. It outputs an edge map and a line map.
2. Simple Structure Upsampler (SSU) §3.2 app. D. A small CNN enlarges the grayscale sketch. The authors use it on the sketch rather than on the RGB image.
3. Structure Feature Encoder (SFE) §3.3. A gated-convolution encoder turns the sketch into four feature maps.
4. Masking positional encoding (MPE) §3.3. A distance map and a direction map computed from the mask, both encoded and added to the painter's first layer.
5. ZeroRA injection §3.3. Each SFE feature is scaled by a zero-initialised scalar and added to a FTR layer input, so the painter's output is unchanged at the start of finetuning.
6. FTR, the texture painter §3.3. It uses the LaMa architecture (FFC blocks, Fourier convolutions). Its losses are in §3.4.

Two things at inference app. E.4: TSR runs Mask-Predict for 5 rounds, re-masking only the low-confidence edge and line pixels, never the input image. Then SSU and the painter run once per image.

Pseudocode (a summary of §3; the paper has no numbered algorithm):

# Inference pipeline, summarised from §3 (not a numbered algorithm in the paper)
# Input:  masked_image (B, 3, 256, 256) in [0, 1]; mask (B, 1, 256, 256), 1 = hole
# Output: inpainted_image (B, 3, 256, 256)

def inpaint(masked_image, mask, edges, lines):
    # 1. Transformer restores edge and line logits  [§3.1]
    edge_logits, line_logits = tsr_forward(masked_image, edges, lines, mask)
    # 2. Mask-Predict: 5 rounds, re-mask low-confidence structure  [app. E.4]
    edge_map, line_map = mask_predict(edge_logits, line_logits, mask, rounds=5)
    # 3. Upsample the sketch with the learned SSU (2x per call)  [§3.2] [eq. 12]
    edge_big, line_big = ssu_upsample(edge_map, line_map)
    # 4. Four coarse-to-fine feature maps from the encoder  [§3.3] [eq. 3]
    structure_feats = sfe_encode(edge_big, line_big, mask_big)
    # 5. Distance and direction from the mask  [§3.3] [eq. 4] [eq. 5]
    distance_emb, direction_emb = masking_position_encoding(mask_big)
    # 6. Add the MPE to the painter's first layer  [§3.3]
    first = conv1(concat(masked_image, mask_big)) + distance_emb + direction_emb
    # 7. Each feature enters via its zero-initialised scalar  [§3.3] [eq. 7]
    x = first
    for k in range(4):
        x = encoder_or_middle_layer_k(x + alpha[k] * structure_feats[k])
    return decode(x)

Key equations (numbering as in the compiled paper; equations 1–11 are in the main text, 12 in appendix D).

TSR attention scores, with row and column axial attention and learned relative position bias §3.1 eq. 1:

\mathbf{A}_{i,j}^{row}=\mathbf{x}_{ri}\mathbf{W}_{rq}\mathbf{W}_{rk}^T\mathbf{x}_{rj}^T+R_{i,j}^{row}, \qquad \mathbf{A}_{i,j}^{col}=\mathbf{x}_{ci}\mathbf{W}_{cq}\mathbf{W}_{ck}^T\mathbf{x}_{cj}^T+R_{i,j}^{col} eq. 1

TSR loss, binary cross-entropy on the predicted continuous sketches §3.1 eq. 2:

\mathcal{L}_{e}=\mathrm{BCE}(\tilde{\mathbf{I}}_e, \hat{\mathbf{I}}_e),\quad \mathcal{L}_{l}=\mathrm{BCE}(\tilde{\mathbf{I}}_l, \hat{\mathbf{I}}_l) eq. 2

SFE features, four maps, the mask resized to the sketch size §3.3 eq. 3:

\mathbf{S}_0,\mathbf{S}_1,\mathbf{S}_2,\mathbf{S}_3=\mathrm{SFE}(\tilde{\mathbf{I}}_{e},\tilde{\mathbf{I}}_{l},\mathbf{M}) eq. 3

MPE distance term, sinusoidal in the clipped distance, d = 64, D_max = 128 §3.3 eq. 4:

\mathbf{P}_{dis,2i}=\sin\!\left(\frac{\mathrm{clip}(\mathbf{D}_{dis},0,D_{max})}{10000^{i/d}}\right),\quad \mathbf{P}_{dis,2i+1}=\cos\!\left(\frac{\mathrm{clip}(\mathbf{D}_{dis},0,D_{max})}{10000^{i/d}}\right) eq. 4

MPE direction term, a learned projection of a 4-channel multi-hot vector §3.3 eq. 5:

\mathbf{P}_{dir}=\mathbf{D}_{dir}\times\mathbf{W}_{dir}\in\mathbb{R}^{256\times256\times d}, \quad \mathbf{W}_{dir}\in\mathbb{R}^{4\times d} eq. 5

ZeroRA, a residual with a zero-initialised trainable weight §3.3 eq. 6:

x'=x+\alpha\cdot F(x) eq. 6

ZeroRA as used to inject structure into the painter §3.3 eq. 7:

\mathbf{X}_{k+1}=\mathrm{Conv}(\mathbf{X}_k+\alpha_k\cdot\mathbf{S}_k),\quad \mathbf{X}_{k+1}=\mathrm{ReLU}(\mathrm{BatchNorm}(\mathbf{X}_{k+1})) eq. 7

Known-region L1 loss, with mask M = 1 for holes §3.4 eq. 8:

\mathcal{L}_{L1}=(1-\mathbf{M})\odot|\hat{\mathbf{I}}-\tilde{\mathbf{I}}|_1 eq. 8

Adversarial loss, as printed §3.4 eq. 9:

\mathcal{L}_{D}=-\mathbb{E}_{\hat{\mathbf{I}}}[\log D(\hat{\mathbf{I}})]-\mathbb{E}_{\tilde{\mathbf{I}},\mathbf{M}}[\log D(\tilde{\mathbf{I}})\odot(1-\mathbf{M})]-\mathbb{E}_{\tilde{\mathbf{I}},\mathbf{M}}[\log(1-D(\tilde{\mathbf{I}}))\odot\mathbf{M}]

\mathcal{L}_{G}=-\mathbb{E}_{\tilde{\mathbf{I}}}[\log D(\tilde{\mathbf{I}})],\quad \mathcal{L}_{adv}=\mathcal{L}_D+\mathcal{L}_G+\lambda_{GP}\mathcal{L}_{GP}

\mathcal{L}_{GP}=\mathbb{E}_{\hat{\mathbf{I}}}\|\nabla_{\hat{\mathbf{I}}}D(\hat{\mathbf{I}})\|^2,\quad \lambda_{GP}=10^{-3} eq. 9

Feature-matching loss L_fm is an L1 loss between discriminator features of real and fake samples; no equation is given §3.4.

High-receptive-field perceptual loss, with a pretrained dilated ResNet-50 segmentation network φ_hrf §3.4 eq. 10:

\mathcal{L}_{hrf}=\mathbb{E}\left(\left[\phi_{hrf}(\hat{\mathbf{I}})-\phi_{hrf}(\tilde{\mathbf{I}})\right]^2\right) eq. 10

Total generator loss in the incremental stage §3.4 eq. 11:

\mathcal{L}_{final}=\lambda_{L1}\mathcal{L}_{L1}+\lambda_{adv}\mathcal{L}_{adv}+\lambda_{fm}\mathcal{L}_{fm}+\lambda_{hrf}\mathcal{L}_{hrf}

with λ_L1 = 10, λ_adv = 10, λ_fm = 100, λ_hrf = 30 §3.4 eq. 11.

SSU output, shifted and scaled sigmoid app. D eq. 12:

\mathbf{I}'_e=\mathrm{sigmoid}[\gamma(\mathbf{I}'_e+\beta)],\quad \mathbf{I}'_l=\mathrm{sigmoid}[\gamma(\mathbf{I}'_l+\beta)] eq. 12

with γ = β = 2 at evaluation and γ, β drawn from [1.5, 3] for finetuning app. D. Applying SSU twice doubles the size again, and the result is resized bilinearly to the target size app. D.

Hyperparameters

Values are the paper's unless marked "code". Sources: §, tab., app., fig., or code citations.

Results

Author-reported. Setup: 256 px images, mixed irregular and segmentation masks, for Places2 and Indoor (own collection, see Pitfalls); baselines are EC, HiFill, MST, Co-Mod and LaMa. The authors retrained EC, MST and LaMa on Indoor §4.3.

Table 1 tab. 1 (256 px, mixed masks):

Table 2 tab. 2 (structure F1, edges and lines; P = precision, R = recall, all in %):

The authors state that these TSR numbers use Mask-Predict, which improves the structure output but not these metrics §4.4.

Table 3 tab. 3 (512 px Places2, finetuned from 256 to 512 px, with and without MPE):

Table 4 tab. 4 (Indoor, 256 px, component ablation; yes = present):

Table 5 tab. 5 (512 and 1024 px):

Table 7 app. E.2 (256 px, by mask ratio; PSNR, SSIM, FID, LPIPS; Indoor columns EC, MST, LaMa, this paper; Places2 columns EC, HiFill, Co-Mod, MST, LaMa, this paper):

The Places2 Mixed LPIPS row of Table 7 conflicts with Table 1 (see Pitfalls).

Table 8 app. E.3 (inference cost, single V100, batch 8):

Table 9 app. E.3 (TSR ablation, Indoor, F1 %):

Table 10 app. E.6 (1k Places2 at 512 px, 40% centre rectangles, no retraining):

Table 11 app. E.7 (512 px texture images, LaMa vs this paper):

Headline (author-reported, 256 px, mixed masks): PSNR 25.57 vs 25.20 for LaMa on Indoor, and 24.42 vs 24.37 on Places2 tab. 1. The authors also report gains over LaMa at 512 and 1024 px tab. 5 §4.7, and a user preference study on Indoor and Places2 app. E.5.

Paper vs code

Official code: DQiaole/ZITS_inpainting at commit 6f4a47bb03bfb82a4860fc55ae0719ab0db8f6eb, linked in the abstract. The owner matches the first author, and the code is Apache-2.0. I did not look for other repositories because this one is linked from the paper.

Differences between paper and code, each with both citations:

• TSR optimiser and schedule. Paper: Adam, lr 6e-4, 1,000 warmup, cosine §4.2. Code: AdamW TSR_model.py, lr 4.24e-4 default TSR_train.py, warmup 1,500 TSR_train.py, betas (0.9, 0.95), weight decay 0 TSR_train.py, cosine floor 0.1 TSR_trainer.py.
• TSR training length. Paper: 150k (Indoor), 400k (Places2) steps §4.2. Code: epochs, default 12 TSR_train.py. The shipped configs do not fix a step count for TSR.
• TSR batch. Paper: 30 app. C. Code default 32 TSR_train.py. The 256 px FTR config uses 32 config_ZITS_places2.yml.
• Loss masking in TSR. Paper: eq. 2 is a plain BCE on the sketches §3.1 eq. 2. Code multiplies the loss by the hole mask and then takes the mean over all pixels, so known pixels contribute zeros TSR_model.py. The input is also multiplied by (1 - mask) TSR_model.py.
• Causal mask. The code defines a causal mask buffer transformer.py, but the block's forward pass never applies it transformer.py. Attention in the TSR is bidirectional. The paper does not mention a causal mask.
• Axial block. The paper defines only the attention scores §3.1 eq. 1. The code uses pre-norm LayerNorm with eps 1e-4 transformer.py, sums row, column and residual outputs transformer.py, then applies a feed-forward layer transformer.py.
• ZeroRA and LayerNorm. The paper says ZeroRA replaces LayerNorm in the transformer §3.3. The released TSR keeps LayerNorm transformer.py. ZeroRA is used only to inject SFE features into the painter LaMa.py.
• MPE gating. The paper adds the distance and direction embeddings to the first layer §3.3 eq. 4 eq. 5. The code multiplies both by zero-initialised scalars alpha5 and alpha6 LaMa.py LaMa.py LaMa.py before adding them to the first convolution output LaMa.py. The MPE therefore contributes nothing at the start of finetuning. The paper's MPE ablation says it was finetuned "with ZeroRA" §4.6, which agrees with the gating.
• Sinusoid layout. Paper: interleaved sine and cosine, even and odd channels §3.3 eq. 4. Code: sine in the first half of the channels and cosine in the second LaMa.py. Code also looks up a frozen table of 128 entries by integer distance LaMa.py LaMa.py, rather than evaluating eq. 4 on a continuous distance.
• Distance clip. Paper: D_max = 128 §3.3. Code: index clipped to 127 dataset_FTR.py.
• Resizing the MPE. Paper: the positional embedding is resized by nearest interpolation across scales §3.3. Code: distances are computed on the 256 px mask, then the index map is nearest-resized to the original image size dataset_FTR.py.
• Canny sigma. Paper: sigma 2.0 at 256 px and 2.5 at 512 px fig. 3. Code defaults to those values FTR_trainer.py, but the Places2 config sets 3.0 and 3.75 config_ZITS_places2.yml, interpolated linearly with image size dataset_FTR.py.
• SSU. Paper: a simple 4-layer CNN that doubles the size §3.2. Code: bilinear 2x, then four convolutions, returning the 2x output upsample.py upsample.py. A second head is defined but unused by the callers upsample.py FTR_trainer.py.
• SSU training perturbation. Random shift and scale in [1.5, 3] is applied only when fix_256 is off FTR_trainer.py. The Places2 256 px config sets fix_256: True config_ZITS_places2.yml, so that config skips the SSU and the perturbation.
• Structure substitution during finetuning. The paper does not describe this. The Places2 config replaces the ground-truth edges and lines with TSR predictions gradually from 808k steps, and fully from 820k up to 850k config_ZITS_places2.yml FTR_trainer.py. During training the number of Mask-Predict rounds is drawn at random from 1 to 5 FTR_trainer.py; inference uses 5 TSR_inference.py.
• Edge thresholding. Paper: no thresholding described; Mask-Predict re-masks low-confidence pixels app. E.4. Code: after each round, edges accumulate and are then binarised at 0.25 utils.py. Line maps stay continuous utils.py. The ranking adds 0.5 to the positive edge class utils.py.
• Wireframes. Paper: lines come from a masking-augmented wireframe detector from Cao et al. §3.1. Code: a separate script runs LSM-HAWP with threshold 0.8 at size 512 lsm_hawp_inference.py. The masking augmentation is not in the repository.
• Loss weights and FTR learning rates. These match: the config lists the same weights and learning rates as the paper config_ZITS_places2.yml config_LAMA.yml.
• Learning-rate decay. Not described in the paper. Code: StepLR x0.5 every 400k steps at 256 px FTR_model.py and every 40k at 512 px config_ZITS_HR_places2.yml.
• Pretrained weights. The finetuning configs load LaMa Places2 checkpoints, plus TSR and SSU checkpoints, from ./ckpt/ config_ZITS_places2.yml. These files are not in the repository, so stage 1 must be trained or downloaded first.

Pitfalls

• Structure at 256 px. The sketch stage works only at 256 px; the authors say distant urban views lose structure and blur app. G fig. 24. Upsampling is learned, but larger scenes still depend on the 256 px sketch.
• Mask-Predict cost. Five TSR forward passes per image at inference app. E.4 TSR_inference.py. Table 8 gives FPS for TSR alone, not the full pipeline app. E.3.
• Zero initialisation delays the effect. Injected features and the MPE start at zero, so the output equals the pretrained painter at the start. The authors report that without ZeroRA the finetune is unstable §4.6 fig. 8.
• Table 1 and Table 7 disagree. For Places2 mixed masks, LPIPS for HiFill, Co-Mod and MST is 0.246, 0.122 and 0.137 in Table 1, but 0.137, 0.246 and 0.122 in Table 7 tab. 1 app. E.2. The other values match. Check the source before quoting these three numbers.
• Dataset. Indoor is the authors' own collection: 20,055 training images (5,000 from ShanghaiTech, 15,055 from NYUDepthV2) and 1,000 validation images §4.1. The collection is not in the repository, so Indoor results are hard to reproduce.
• Stage 1 needs a pretrained inpainter. The incremental method starts from a pretrained LaMa-style painter §3.3. Stage 1 is 800k steps on Places2 §4.2.
• Discriminator loss as printed. Eq. 9 places element-wise masks inside the log terms. The paper does not say how to reduce the masked terms to a scalar. The code routes the mask through mask_as_fake_target config_ZITS_places2.yml. Check that before reimplementing.
• L1 normalisation. Eq. 8 is printed without a mean or sum. The paper does not say which is used.
• Feature-matching loss. L_fm has no equation in the paper §3.4.
• Perceptual network. The HRF loss needs a dilated ResNet-50 segmentation network, which the paper describes but does not provide §3.4. The config points to a LaMa ADE20K file at ./ config_ZITS_places2.yml, not in the repository.
• NumPy alias. The dataset loader uses np.float dataset_FTR.py. Recent NumPy releases no longer provide that alias, so use float when reimplementing.
• Mask ratio mapping. The paper's "20%" combined-mask wording §4.2 and the config's MASK_RATE thresholds config_ZITS_places2.yml are not reconciled in the paper. I did not trace the mapping in datasets/dataset_FTR.py.
• Two resolutions, two pipelines. The Places2 256 px config bypasses the SSU config_ZITS_places2.yml. The 512 px config does not set fix_256 and uses the SSU config_ZITS_HR_places2.yml. Results in Table 5 come from the second pipeline.

Lineage

• Builds on: LaMa (Suvorov et al., 2021) as the FTR painter and loss source §3.3 §3.4; MST (Cao et al.) for the sketch space and the structure baseline §2 §4.3; ReZero (Bachlechner et al., 2020) for zero-initialised residuals §3.3; axial attention (Ho et al., 2019) and relative position encoding (Raffel et al., 2019) in the TSR §3.1; gated convolutions (Yu et al., 2019) in the SFE §3.3.
• Compared against: EC (Nazeri et al.), HiFill (Yi et al.), MST, Co-Mod (Zhao et al.) and LaMa §4.3.
• Superseded by: not known from the content of this wiki.