U-Net for Image Inpainting on CelebA-HQ (256x256)

This repository contains a PyTorch implementation of a deep U-Net with Residual Blocks, trained to perform image inpainting on the high-resolution CelebA-HQ dataset. The model takes an image with a masked (blacked-out) region and intelligently reconstructs the missing content.

Model Description

The model uses a ComplexUNet architecture, which is a variant of the standard U-Net adapted for higher-resolution images.

Deeper Architecture: Features 5 downsampling and 5 upsampling stages to effectively handle the 256x256 resolution.
Residual Blocks: Incorporates residual blocks in each stage for more stable training and to combat vanishing gradients in a deep network.
Width: The model was trained with base_channels=64.
Total Parameters: 129,973,507

How to Use

First, ensure you have the required libraries installed:

pip install torch torchvision numpy Pillow

Next, download the model weights file (inpainting_model_celebahq.pth) from the "Files and versions" tab of this repository. Then, you can use the following code to load the model and perform inpainting.

import torch
from torchvision import transforms as T
from PIL import Image
from model import ComplexUNet # Import the class from the model.py file in this repo

# --- 1. Setup ---
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
MODEL_PATH = "inpainting_model_celebahq.pth" # Path to your downloaded weights

# --- 2. Load Model ---
model = ComplexUNet(base_channels=64)
model.load_state_dict(torch.load(MODEL_PATH, map_location=DEVICE))
model.to(DEVICE)
model.eval()

# --- 3. Load and Preprocess Image ---
# As an example, we create a dummy tensor. Replace this with your own image.
# image = Image.open("your_image.png").convert("RGB")
# transform = T.Compose([T.Resize((256, 256)), T.ToTensor()])
# image_tensor = transform(image)
image_tensor = torch.rand(3, 256, 256) # Dummy image tensor

# --- 4. Create a Mask ---
# This creates a black square in the center of the image.
masked_tensor = image_tensor.clone()
masked_tensor[:, 64:192, 64:192] = 0

# --- 5. Perform Inpainting ---
with torch.no_grad():
    input_tensor = masked_tensor.unsqueeze(0).to(DEVICE)
    reconstructed_tensor = model(input_tensor).squeeze(0).cpu()

# --- 6. Save the Result ---
from torchvision.transforms.functional import to_pil_image
reconstructed_image = to_pil_image(reconstructed_tensor)
reconstructed_image.save("reconstructed_image.png")
print("✅ Inpainting complete. Saved to reconstructed_image.png")

Training & Evaluation

Training Data

The model was trained on the CelebA-HQ dataset, sourced from Kaggle (lamsimon/celebahq).

Preprocessing: Images were resized to 256x256 pixels.
Augmentation: During training, a random rectangular mask was applied to each image to create the input for the inpainting task.

Training Procedure

Framework: PyTorch
Optimizer: Adam
Learning Rate: 1e-3
Epochs: 20
Batch Size: 16
Loss Function: Mean Squared Error (MSE)
Final Training Loss: 0.001845

Evaluation Results

The model was evaluated on a held-out test set of 3,000 images from the CelebA-HQ dataset, achieving the following metrics:

Metric	Value
PSNR	29.03 dB
SSIM	0.950

tahamajs/celebahq-inpainting-unet