GAN that tries to create a rainbow. (For Brex's Rainbow campaign: https://www.facebook.com/BrexHQ/posts/932319587222145)

brexrainbowgan.py

# -*- coding: utf-8 -*-
"""BrexRainbowGAN.ipynb

Author: Anoop Dixith
For Brex Rainbow initiative: https://www.facebook.com/BrexHQ/posts/932319587222145
"""

import tensorflow as tf

import glob
import imageio
import matplotlib.pyplot as plt
import numpy as np
import os
import PIL
from tensorflow.keras import layers
import time
from google.colab import files
from IPython.display import display
from IPython import display

uploaded = files.upload()


def decode_img(img):
    img = tf.image.decode_jpeg(img, channels=3)
    img = tf.image.convert_image_dtype(img, tf.float32)
    return tf.image.resize(img, [100, 100])


img1 = decode_img(tf.io.read_file('rb1.png'))
img2 = decode_img(tf.io.read_file('rb2.png'))
img3 = decode_img(tf.io.read_file('rb3.png'))
img4 = decode_img(tf.io.read_file('rb4.png'))
img5 = decode_img(tf.io.read_file('rb5.png'))
img6 = decode_img(tf.io.read_file('rb6.png'))
img7 = decode_img(tf.io.read_file('rb7.png'))
img8 = decode_img(tf.io.read_file('rb8.png'))
img9 = decode_img(tf.io.read_file('rb9.png'))
img10 = decode_img(tf.io.read_file('rb10.png'))
img11 = decode_img(tf.io.read_file('rb11.png'))
img12 = decode_img(tf.io.read_file('rb12.png'))
img13 = decode_img(tf.io.read_file('rb13.png'))
img14 = decode_img(tf.io.read_file('rb14.png'))
img15 = decode_img(tf.io.read_file('rb15.png'))
img16 = decode_img(tf.io.read_file('rb16.png'))
img17 = decode_img(tf.io.read_file('rb17.png'))
img18 = decode_img(tf.io.read_file('rb18.png'))
img19 = decode_img(tf.io.read_file('rb19.png'))
img20 = decode_img(tf.io.read_file('rb20.png'))
img21 = decode_img(tf.io.read_file('rb21.png'))
img22 = decode_img(tf.io.read_file('rb22.png'))
img23 = decode_img(tf.io.read_file('rb23.png'))
img24 = decode_img(tf.io.read_file('rb24.png'))
img25 = decode_img(tf.io.read_file('rb25.png'))
img26 = decode_img(tf.io.read_file('rb26.png'))
img27 = decode_img(tf.io.read_file('rb27.png'))
img28 = decode_img(tf.io.read_file('rb28.png'))
img29 = decode_img(tf.io.read_file('rb29.png'))
img30 = decode_img(tf.io.read_file('rb30.png'))
img31 = decode_img(tf.io.read_file('rb31.png'))
img32 = decode_img(tf.io.read_file('rb32.png'))
img33 = decode_img(tf.io.read_file('rb33.png'))
img34 = decode_img(tf.io.read_file('rb34.png'))
img35 = decode_img(tf.io.read_file('rb35.png'))
img36 = decode_img(tf.io.read_file('rb36.png'))
img37 = decode_img(tf.io.read_file('rb37.png'))
img38 = decode_img(tf.io.read_file('rb38.png'))
img39 = decode_img(tf.io.read_file('rb39.png'))
img40 = decode_img(tf.io.read_file('rb40.png'))
img41 = decode_img(tf.io.read_file('rb41.png'))
img42 = decode_img(tf.io.read_file('rb42.png'))
img43 = decode_img(tf.io.read_file('rb43.png'))
img44 = decode_img(tf.io.read_file('rb44.png'))
img45 = decode_img(tf.io.read_file('rb45.png'))
img46 = decode_img(tf.io.read_file('rb46.png'))
img47 = decode_img(tf.io.read_file('rb47.png'))
img48 = decode_img(tf.io.read_file('rb48.png'))
img49 = decode_img(tf.io.read_file('rb49.png'))
img50 = decode_img(tf.io.read_file('rb50.png'))

all_images = [img1, img2, img3, img4, img5, img6, img7, img8, img9, img10, img11, img12, img13, img14, img15, img16,
              img17, img18, img19, img20, img21, img22, img23, img24, img25, img26, img27, img28, img29, img30, img31,
              img32, img33, img34, img35, img36, img37, img38, img39, img40, img41, img42, img43, img44, img45, img46,
              img47, img48, img49, img50]
all_labels = tf.constant(
    [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
     1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1])

for each in all_images:
    print("Image shape: ", each.numpy().shape)

train_images = all_images
train_labels = all_labels
# Normalize
train_images = (np.asarray(train_images, dtype=np.float32) + (-127.5)) / 127.5
print(train_images)
BUFFER_SIZE = 50
BATCH_SIZE = 8
train_dataset = tf.data.Dataset.from_tensor_slices(train_images).shuffle(BUFFER_SIZE).batch(BATCH_SIZE)


def make_generator_model():
    model = tf.keras.Sequential()
    model.add(layers.Dense(25 * 25 * 256, use_bias=False, input_shape=(100,)))
    model.add(layers.BatchNormalization())
    model.add(layers.LeakyReLU())

    model.add(layers.Reshape((25, 25, 256)))
    assert model.output_shape == (None, 25, 25, 256)  # Note: None is the batch size

    model.add(layers.Conv2DTranspose(128, (25, 25), strides=(1, 1), padding='same', use_bias=False))
    print(model.output_shape)
    assert model.output_shape == (None, 25, 25, 128)
    model.add(layers.BatchNormalization())
    model.add(layers.LeakyReLU())

    model.add(layers.Conv2DTranspose(64, (50, 50), strides=(2, 2), padding='same', use_bias=False))
    assert model.output_shape == (None, 50, 50, 64)
    model.add(layers.BatchNormalization())
    model.add(layers.LeakyReLU())

    model.add(layers.Conv2DTranspose(64, (50, 50), strides=(2, 2), padding='same', use_bias=False))
    assert model.output_shape == (None, 100, 100, 64)
    model.add(layers.BatchNormalization())
    model.add(layers.LeakyReLU())

    # model.add(layers.Conv2DTranspose(64, (20, 30), strides=(2, 2), padding='same', use_bias=False))
    # assert model.output_shape == (None, 40, 60, 64)
    # model.add(layers.BatchNormalization())
    # model.add(layers.LeakyReLU())

    model.add(layers.Conv2DTranspose(3, (100, 100), strides=(1, 1), padding='same', use_bias=False, activation='tanh'))
    assert model.output_shape == (None, 100, 100, 3)

    return model


generator = make_generator_model()

noise = tf.random.normal([1, 100])
generated_image = generator(noise, training=False)

plt.imshow(generated_image[0, :, :, :])


def make_discriminator_model():
    model = tf.keras.Sequential()
    model.add(layers.Conv2D(64, (5, 5), strides=(2, 2), padding='same',
                            input_shape=[100, 100, 3]))
    model.add(layers.LeakyReLU())
    model.add(layers.Dropout(0.3))

    model.add(layers.Conv2D(128, (5, 5), strides=(2, 2), padding='same'))
    model.add(layers.LeakyReLU())
    model.add(layers.Dropout(0.3))

    model.add(layers.Flatten())
    model.add(layers.Dense(1))

    return model


discriminator = make_discriminator_model()
decision = discriminator(generated_image)
print(decision)

cross_entropy = tf.keras.losses.BinaryCrossentropy(from_logits=True)


def discriminator_loss(real_output, fake_output):
    real_loss = cross_entropy(tf.ones_like(real_output), real_output)
    fake_loss = cross_entropy(tf.zeros_like(fake_output), fake_output)
    total_loss = real_loss + fake_loss
    return total_loss


def generator_loss(fake_output):
    return cross_entropy(tf.ones_like(fake_output), fake_output)


generator_optimizer = tf.keras.optimizers.Adam(1e-4)
discriminator_optimizer = tf.keras.optimizers.Adam(1e-4)

checkpoint_dir = './rainbow_checkpoints'
checkpoint_prefix = os.path.join(checkpoint_dir, "ckpt")
checkpoint = tf.train.Checkpoint(generator_optimizer=generator_optimizer,
                                 discriminator_optimizer=discriminator_optimizer,
                                 generator=generator,
                                 discriminator=discriminator)

EPOCHS = 25
noise_dim = 100
num_examples_to_generate = 10

# We will reuse this seed overtime (so it's easier)
# to visualize progress in the animated GIF)
seed = tf.random.normal([num_examples_to_generate, noise_dim])


@tf.function
def train_step(images):
    noise = tf.random.normal([BATCH_SIZE, noise_dim])

    with tf.GradientTape() as gen_tape, tf.GradientTape() as disc_tape:
        generated_images = generator(noise, training=True)

        real_output = discriminator(images, training=True)
        fake_output = discriminator(generated_images, training=True)

        gen_loss = generator_loss(fake_output)
        disc_loss = discriminator_loss(real_output, fake_output)

    gradients_of_generator = gen_tape.gradient(gen_loss, generator.trainable_variables)
    gradients_of_discriminator = disc_tape.gradient(disc_loss, discriminator.trainable_variables)

    generator_optimizer.apply_gradients(zip(gradients_of_generator, generator.trainable_variables))
    discriminator_optimizer.apply_gradients(zip(gradients_of_discriminator, discriminator.trainable_variables))


def train(dataset, epochs):
    for epoch in range(epochs):
        start = time.time()

        for image_batch in dataset:
            train_step(image_batch)

        generate_and_save_images(generator,
                                 epoch + 1,
                                 seed)

        if (epoch + 1) % 15 == 0:
            checkpoint.save(file_prefix=checkpoint_prefix)

        print('Time for epoch {} is {} sec'.format(epoch + 1, time.time() - start))

    # Generate after the final epoch
    # display.clear_output(wait=True)
    generate_and_save_images(generator,
                             epochs,
                             seed)


def generate_and_save_images(model, epoch, test_input):
    predictions = model(test_input, training=False)

    fig = plt.figure(figsize=(25, 25))
    for i in range(predictions.shape[0]):
        plt.subplot(10, 10, i + 1)
        plt.imshow(predictions[i, :, :, 0] * 127.5 + 127.5, cmap='Accent')
        plt.axis('off')
        plt.title('1')

    plt.savefig('image_at_epoch_{:04d}.png'.format(epoch))
    plt.show()

    fig = plt.figure(figsize=(10, 10))
    for i in range(predictions.shape[0]):
        plt.subplot(10, 10, i + 1)
        plt.imshow(predictions[i, :, :, 0] * 127.5 + 127.5, cmap='Accent')
        plt.axis('off')
        plt.title('2')

    plt.savefig('image_at_epoch_{:04d}.png'.format(epoch))
    plt.show()

    fig = plt.figure(figsize=(25, 25))
    for i in range(predictions.shape[0]):
        plt.subplot(10, 10, i + 1)
        plt.imshow(predictions[i, :, :, :] * 127.5 + 127.5, cmap='Accent')
        plt.axis('off')
        plt.title('3')

    plt.savefig('image_at_epoch_{:04d}.png'.format(epoch))
    plt.show()

    fig = plt.figure(figsize=(10, 10))
    for i in range(predictions.shape[0]):
        plt.subplot(10, 10, i + 1)
        plt.imshow(predictions[i, :, :, :] * 127.5 + 127.5, cmap='Accent')
        plt.axis('off')
        plt.title('4')

    plt.savefig('image_at_epoch_{:04d}.png'.format(epoch))
    plt.show()

train(train_dataset, EPOCHS)