DOCSEARCHbyAlgolia
Submit Docs

DeOldify

repository·master·Indexed 12 days ago

https://github.com/jantic/deoldify

A deep-learning project that uses Generative Adversarial Network (GAN) techniques and NoGAN training to colorize and restore old black-and-white images and film footage. It offers Artistic, Stable, and Video models with different backbones (resnet34 and resnet101) to balance vibrance and temporal stability.

Tokens
14.9K
Snippets
42
Records
58
Agent score
93%

What's inside DeOldify

  1. Understand NoGAN training

    master

    NoGAN is a training technique designed to provide the realism of Generative Adversarial Networks (GANs) while minimizing the artifacts (glitches, flickering, and color oscillations) typically associated with them.

    The NoGAN Workflow:

    1. Generator Pretraining: Train the generator using conventional methods with only feature loss.
    2. Critic Pretraining: Generate images from the pre-trained generator and train a critic to distinguish them from real images (binary classification).
    3. GAN Training: Train the generator and critic together in a GAN setting at the target resolution (e.g., 192px).

    The Inflection Point: Productive GAN training occurs within a very small window. There is an 'inflection point' where the critic has transferred all useful information to the generator. Beyond this point, image quality tends to oscillate or degrade (e.g., skin turning orange or lips becoming overly red). Successful NoGAN implementation requires identifying this point to avoid overtraining.

  3. Download and configure pretrained weights

    master

    To perform colorization inference without training your own models, you must download the completed weights and place them in the /models/ directory of your DeOldify installation.

    There are three main model types: Artistic, Stable, and Video. For inference, you typically need both the Generator and the Critic weights.

    ### Completed Generator Weights
- [Artistic](https://data.deepai.org/deoldify/ColorizeArtistic_gen.pth)
- [Stable](https://www.dropbox.com/s/axsd2g85uyixaho/ColorizeStable_gen.pth?dl=0)
- [Video](https://data.deepai.org/deoldify/ColorizeVideo_gen.pth)

### Completed Critic Weights
- [Artistic](https://www.dropbox.com/s/xpq2ip9occuzgen/ColorizeArtistic_crit.pth?dl=0)
- [Stable](https://www.dropbox.com/s/s53699e9n84q6sp/ColorizeStable_crit.pth?dl=0)
- [Video](https://www.dropbox.com/s/xnq1z1oppvgpgtn/ColorizeVideo_crit.pth?dl=0)

  4. Install DeOldify locally via Anaconda

    master

    To run DeOldify on your own machine, you can use Anaconda to manage the environment. This method installs the necessary dependencies and sets up a Jupyter Lab environment.

    Requirements:

    • OS: Linux (Ubuntu 18.04 or 16.04 recommended). Windows is not supported.
    • Hardware (Colorization only): A decent graphics card with approximately 4GB+ VRAM.
    • Hardware (Training only): A high-end graphics card (e.g., 11GB+ VRAM recommended).
    git clone https://github.com/jantic/DeOldify.git DeOldify
cd DeOldify
conda env create -f environment.yml

source activate deoldify
# Or if using a recent conda version:
# conda activate deoldify

jupyter lab

  6. Achieve stable video colorization

    master

    To achieve stable, flicker-free video colorization in DeOldify, the system uses NoGAN training which avoids temporal modeling in favor of high-quality isolated image generation. To improve stability when rendering video, consider the following:

    • Increase render_factor: Rendering at a higher resolution (higher render_factor) provides the model with higher fidelity information, leading to more consistent colorization decisions.
    • Use the Stable model architecture: The 'Stable' model uses a resnet101 backbone instead of resnet34, which allows for more consistent object detection (e.g., skin rendering) and reduces artifacts like 'zombie hands'.
    • NoGAN Training: DeOldify utilizes NoGAN training to combine the vibrant colorization of GANs with the stability of conventional methods, minimizing the flickering artifacts common in standard GANs.

  7. Use DeOldify via Google Colab

    master

    For running DeOldify in a cloud environment with GPU support, use the official Google Colab notebooks. There are two main model types available:

    1. Artistic Model

    Produces more interesting and vibrant colors but may have more glitches.

    2. Stable Model

    Produces less 'interesting' colors but significantly reduces visual glitches and artifacts.

    Tutorial: A video tutorial for using these Colabs is available here.

  10. Use inference notebooks for colorization

    master

    Once the pretrained weights are placed in the /models/ folder, you can use the following Jupyter notebooks to guide your colorization tasks:

    • ImageColorizerArtistic.ipynb (for artistic style images)
    • ImageColorizerStable.ipynb (for stable style images)
    • VideoColorizer.ipynb (for video colorization)

  11. Use custom samplers to iterate through large datasets

    master

    Because the ImageNet dataset is extremely large, iterating through every sample in every epoch is impractical. This notebook uses a FixedLenRandomSampler to limit the number of samples per epoch while still slowly traversing the entire dataset over many epochs. This allows for more frequent validation and metric logging.

    To implement this, the ImageDataBunch.create method is monkey-patched to accept a sampler argument, and get_data is used to wrap the FixedLenRandomSampler into the data loading pipeline.

    # Custom sampler to limit epoch size
class FixedLenRandomSampler(RandomSampler):
    def __init__(self, data_source, epoch_size):
        super().__init__(data_source)
        self.epoch_size = epoch_size
        self.not_sampled = np.array([True]*len(data_source))

    @property
    def reset_state(self): self.not_sampled[:] = True

    def __iter__(self):
        # Logic to sample a fixed number of items per epoch
        ...
        return iter(idx)

    def __len__(self):
        return self.epoch_size

# Usage in data loading
def get_data(bs:int, sz:int, keep_pct=1.0, random_seed=None, valid_pct=0.2, epoch_size=1000):
    train_sampler = partial(FixedLenRandomSampler, epoch_size=epoch_size)
    samplers = [train_sampler, SequentialSampler, SequentialSampler, SequentialSampler]
    return get_colorize_data(sz=sz, bs=bs, crappy_path=path_lr, good_path=path_hr, 
                             random_seed=random_seed, keep_pct=keep_pct, 
                             samplers=samplers, valid_pct=valid_pct)

  12. The Repeatable GAN Training Cycle

    master

    After pre-training the generator, you enter a GAN cycle. For best results, repeat this cycle approximately 5-8 times. In each cycle, you increment the checkpoint numbers to avoid overwriting previous progress. The cycle consists of:

    1. Saving Generated Images: Use the current generator to predict and save images to monitor quality.
    2. Pre-training/Fine-tuning the Critic: Train the critic to distinguish between real color images and generated ones.
    3. GAN Training: Train the generator and critic together using GANLearner.
    # Incrementing checkpoint logic for repeatable cycles
old_checkpoint_num = 0
checkpoint_num = old_checkpoint_num + 1
gen_old_checkpoint_name = gen_name + '_' + str(old_checkpoint_num)
gen_new_checkpoint_name = gen_name + '_' + str(checkpoint_num)
crit_old_checkpoint_name = crit_name + '_' + str(old_checkpoint_num)
crit_new_checkpoint_name = crit_name + '_' + str(checkpoint_num)