This project is an implementation of an image inpainting Generative Adversarial Network (GAN) designed to reconstruct missing or masked areas of images. The project utilizes TensorFlow's Keras to build and train a GAN that can generate realistic image patches to fill in masked regions, supporting multiple well-known datasets, including MNIST, CIFAR-10, CelebA, and more.
The architecture consists of two primary components: a Generator and a Discriminator. The Generator is responsible for generating the missing regions, while the Discriminator evaluates the quality of the generated inpainting, thereby improving the Generator through adversarial training.
- Multiple Mask Types: Supports different types of masks, including square, circular, and arbitrary polygonal masks.
- Dataset Flexibility: Can be trained on a wide range of datasets, including MNIST, EMNIST, GTSRB, CIFAR-10, ImageNet, COCO, and CelebA.
- Incremental Saving: Models are saved incrementally after each epoch, allowing for resuming from the latest saved state.
- Inpainting Quality Metrics: Utilizes SSIM (Structural Similarity Index Measure) and PSNR (Peak Signal-to-Noise Ratio) to evaluate the quality of generated inpainted images.
- Customizable Training: Includes adjustable parameters for learning rate, epochs, and image size.
main.py: Entry point for the project. Loads datasets, initializes or loads models, and orchestrates the training process.train.py: Handles the training loop for the GAN, including forward passes for both Generator and Discriminator, calculating losses, and updating model weights.models.py: Contains the architectures for the Generator and Discriminator models, built using TensorFlow and Keras layers.data_loader.py: Loads and preprocesses datasets using TensorFlow Datasets, including resizing, normalizing, and augmenting images.incremental_saver.py: Manages model saving, including retrieving the latest saved model for training continuation and saving checkpoints during training.test.py: Allows users to select an image, apply a mask, and view the inpainting results generated by the trained model.utils.py: Utility functions for creating masks, calculating metrics, and plotting images.Constants.py: Contains hyperparameters and configuration constants used throughout the project (e.g., learning rate, image size, mask ratios).
- Python 3.8+
- TensorFlow 2.10+
- TensorFlow Datasets (tfds)
- NumPy
- OpenCV
- Matplotlib
- Tqdm
Install the dependencies using:
pip install -r requirements.txtTo train the model, run the following command:
python main.py- EPOCHS: Set the number of epochs to train the model (default: 50).
- IMAGE_SIZE: Set the size to which input images will be resized (default: 128).
- BATCH_SIZE: Set the batch size for training (default: 64).
- INITIAL_LEARNING_RATE: Initial learning rate for both Generator and Discriminator.
- DECAY_RATE & DECAY_STEPS: Used for learning rate scheduling.
To test the inpainting capabilities of the trained model, run the test.py script:
python test.pyThe script allows you to select an image, apply a mask, and visualize the result generated by the model, along with extracted feature maps from the generator's encoder layers.
The model generates inpainted images based on masked input images. The training process optimizes several loss functions, including adversarial loss, L1 loss, SSIM, and perceptual loss using a VGG16 model to ensure high-quality inpainting results.
- SSIM: Measures the perceptual difference between the inpainted and original images.
- PSNR: Evaluates the reconstruction quality in terms of pixel-level differences.
The project leverages TensorFlow and TensorFlow Datasets for building and training the models. The design is inspired by classical inpainting methods and generative adversarial networks for filling missing regions of images.
This project is licensed under the MIT License. See the LICENSE file for details.