This repository contains the code for training a CycleGAN model for unpaired image-to-image translation, inspired by the research paper Unpaired Image-to-Image Translation using Cycle-Consistent Adversarial Networks: https://arxiv.org/abs/1703.10593.
Here's a breakdown of the code files:
-
Discriminator.py and Generator.py: These files define the core components of the CycleGAN model:
- Discriminator (Discriminator.py): This network aims to distinguish between real images and images generated by the generators. It takes two images (e.g., horse and zebra) as input and outputs a single value representing the likelihood of being a real image.
- Generator (Generator.py): This network learns to translate images from one domain (e.g., horse) to another (e.g., zebra). It takes an image from one domain as input and generates a realistic image in the target domain.
-
utils.py: This file contains helper functions for training, including saving/loading model checkpoints and setting random seeds for reproducibility.
-
config.py: This file stores the configuration details for training, such as:
- Device to use (CPU or GPU)
- Training data directories
- Hyperparameters like learning rate, batch size, and number of epochs
- Data augmentation techniques
- Options for loading/saving model checkpoints
-
Dataset.py: This file defines a custom dataset class that loads horse and zebra images, applies data augmentation, and prepares them for training.
-
train.ipynb: This Jupyter notebook is the main training script. It:
- Defines the models (Discriminator and Generator)
- Creates optimizers for training
- Defines loss functions for measuring errors
- Loads the training and validation datasets
- Implements the training loop:
- Trains the Discriminators to distinguish real from fake images
- Trains the Generators to fool the Discriminators and generate realistic images
- Calculates cycle consistency loss to ensure images can be translated back to their original domain
- Optionally calculates identity loss to preserve image content during translation
- Saves the model checkpoints periodically
- Optionally visualizes the generated images during training
Overall, this code implements a CycleGAN model for unpaired image-to-image translation. The model learns to translate images from one domain to another while maintaining image consistency through cycle consistency loss.