tahamajs's picture
Update Readme
7fc8ab7 verified
metadata
license: mit
language: en
library_name: pytorch
tags:
  - image-inpainting
  - computer-vision
  - pytorch
  - unet
  - cifar-10
datasets:
  - cifar10

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