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 96 for 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
Downloads last month

-

Downloads are not tracked for this model. How to track
Inference Providers NEW
This model isn't deployed by any Inference Provider. 🙋 Ask for provider support

Dataset used to train tahamajs/cifar10-inpainting-complex-unet