U-Net for Image Inpainting on CIFAR-10
This repository contains a PyTorch implementation of a deep U-Net with Residual Blocks, trained to perform image inpainting on the CIFAR-10 dataset. The model takes a 32x32 image with a masked (blacked-out) region and reconstructs the missing part.
| Original | Masked Input | Reconstructed Output |
|---|---|---|
Note: The images above are illustrative examples. You can generate your own by running the code below.
Model Architecture
The model is a ComplexUNet, a variant of the standard U-Net architecture, designed to be deeper and wider for improved performance.
- Framework: PyTorch
- Architecture: U-Net with 4 downsampling and 4 upsampling stages.
- Backbone: Each stage uses Residual Blocks instead of simple convolutional layers.
- Model Width: The number of base channels is increased to
96for higher capacity. - Total Parameters: 73,148,259
How to Use
The following code snippet provides a complete example of how to load the model, process an image from the CIFAR-10 test set, and visualize the inpainting result.
First, ensure you have the necessary libraries installed:
pip install torch torchvision numpy matplotlib Pillow