Image Completion and Hole Filling with Deep Learning

Module: Computer Vision | Difficulty: Premium

Context Encoder Loss

Partial Convolution

Edge-Aware Loss

where is the Sobel gradient operator.

import torch
import torch.nn as nn

class PartialConv2d(nn.Module):
    def __init__(self, in_c, out_c, kernel=3, stride=1, padding=1):
        super().__init__()
        self.conv = nn.Conv2d(in_c, out_c, kernel, stride, padding, bias=False)
        self.mask_conv = nn.Conv2d(in_c, 1, kernel, stride, padding, bias=False)
        nn.init.constant_(self.mask_conv.weight, 1.0)
        nn.init.constant_(self.mask_conv.bias, 0.0)
        self.bias = nn.Parameter(torch.zeros(out_c))

    def forward(self, x, mask):
        masked_x = x * mask
        output = self.conv(masked_x)
        mask_ratio = self.mask_conv(mask)
        mask_ratio = torch.clamp(mask_ratio, min=1e-8)
        output = output / mask_ratio
        output = output + self.bias.view(1, -1, 1, 1)
        new_mask = (mask_ratio > 0).float()
        return output, new_mask

class InpaintingGenerator(nn.Module):
    def __init__(self):
        super().__init__()
        self.encoder = nn.Sequential(
            PartialConv2d(4, 64, 7, 1, 3), nn.ReLU(inplace=True),
            PartialConv2d(64, 128, 4, 2, 1), nn.ReLU(inplace=True),
            PartialConv2d(128, 256, 4, 2, 1), nn.ReLU(inplace=True),
        )
        self.mid = nn.Sequential(
            nn.Conv2d(256, 256, 3, 1, 1), nn.ReLU(inplace=True),
            nn.Conv2d(256, 256, 3, 1, 1), nn.ReLU(inplace=True),
        )
        self.decoder = nn.Sequential(
            nn.Upsample(scale_factor=2),
            nn.Conv2d(256, 128, 3, 1, 1), nn.ReLU(inplace=True),
            nn.Upsample(scale_factor=2),
            nn.Conv2d(128, 64, 3, 1, 1), nn.ReLU(inplace=True),
            nn.Conv2d(64, 3, 7, 1, 3), nn.Tanh(),
        )

    def forward(self, masked_image, mask):
        x = torch.cat([masked_image * mask, mask], dim=1)
        for layer in self.encoder:
            if isinstance(layer, PartialConv2d):
                x, _ = layer(x, mask)
            else:
                x = layer(x)
        x = self.mid(x)
        return self.decoder(x)

Research Insight: Diffusion-based inpainting models have surpassed GAN-based approaches by generating diverse, coherent completions for large missing regions. The key advantage is that diffusion models naturally handle multimodal outputs (multiple plausible completions) rather than mode-averaging. MAT (Mask-Aware Transformer) uses masked attention to prevent information leakage from known regions, achieving high-quality completion while maintaining global coherence. The challenge remains hole filling in videos where temporal consistency adds complexity.