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 |