ramap27/cyclegan-unpaired-image-translation

CycleGAN implementation for unpaired image-to-image translation. Train a model to translate images between domains (e.g., horse <-> zebra) using PyTorch. This code includes pre-processing, training, and utility functions.

CycleGAN: Unpaired Image-to-Image Translation

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.