cadl.cycle_gan — CADL 1.0.0 documentation

# Source code for cadl.cycle_gan

    """Cycle Generative Adversarial Network for Unpaired Image to Image translation.
    """
    """
    Copyright 2017 Parag K. Mital.  See also NOTICE.md.

    Licensed under the Apache License, Version 2.0 (the "License");
    you may not use this file except in compliance with the License.
    You may obtain a copy of the License at

        http://www.apache.org/licenses/LICENSE-2.0

    Unless required by applicable law or agreed to in writing, software
    distributed under the License is distributed on an "AS IS" BASIS,
    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    See the License for the specific language governing permissions and
    limitations under the License.
    """
    import numpy as np
    import os
    import matplotlib
    matplotlib.use('Agg')
    import matplotlib.pyplot as plt
    import tensorflow as tf
    import tensorflow.contrib.layers as tfl
    from cadl.utils import imcrop_tosquare
    from scipy.misc import imresize

    def l1loss(x, y):
        """Summary

        Parameters
        ----------
        x : TYPE
            Description
        y : TYPE
            Description

        Returns
        -------
        TYPE
            Description
        """
        return tf.reduce_mean(tf.abs(x - y))

    def l2loss(x, y):
        """Summary

        Parameters
        ----------
        x : TYPE
            Description
        y : TYPE
            Description

        Returns
        -------
        TYPE
            Description
        """
        return tf.reduce_mean(tf.squared_difference(x, y))

    def lrelu(x, leak=0.2, name="lrelu"):
        """Summary

        Parameters
        ----------
        x : TYPE
            Description
        leak : float, optional
            Description
        name : str, optional
            Description

        Returns
        -------
        TYPE
            Description
        """
        with tf.variable_scope(name):
            return tf.maximum(x, leak * x)

    def instance_norm(x, epsilon=1e-5):
        """Instance Normalization.

        See Ulyanov, D., Vedaldi, A., & Lempitsky, V. (2016).
        Instance Normalization: The Missing Ingredient for Fast Stylization,
        Retrieved from http://arxiv.org/abs/1607.08022

        Parameters
        ----------
        x : TYPE
            Description
        epsilon : float, optional
            Description

        Returns
        -------
        TYPE
            Description
        """
        with tf.variable_scope('instance_norm'):
            mean, var = tf.nn.moments(x, [1, 2], keep_dims=True)
            scale = tf.get_variable(
                name='scale',
                shape=[x.get_shape()[-1]],
                initializer=tf.truncated_normal_initializer(mean=1.0, stddev=0.02))
            offset = tf.get_variable(
                name='offset',
                shape=[x.get_shape()[-1]],
                initializer=tf.constant_initializer(0.0))
            out = scale * tf.div(x - mean, tf.sqrt(var + epsilon)) + offset
            return out

    def conv2d(inputs,
               activation_fn=lrelu,
               normalizer_fn=instance_norm,
               scope='conv2d',
               **kwargs):
        """Summary

        Parameters
        ----------
        inputs : TYPE
            Description
        activation_fn : TYPE, optional
            Description
        normalizer_fn : TYPE, optional
            Description
        scope : str, optional
            Description
        **kwargs
            Description

        Returns
        -------
        TYPE
            Description
        """
        with tf.variable_scope(scope or 'conv2d'):
            h = tfl.conv2d(
                inputs=inputs,
                activation_fn=None,
                normalizer_fn=None,
                weights_initializer=tf.truncated_normal_initializer(stddev=0.02),
                biases_initializer=None,
                **kwargs)
            if normalizer_fn:
                h = normalizer_fn(h)
            if activation_fn:
                h = activation_fn(h)
            return h

    def conv2d_transpose(inputs,
                         activation_fn=lrelu,
                         normalizer_fn=instance_norm,
                         scope='conv2d_transpose',
                         **kwargs):
        """Summary

        Parameters
        ----------
        inputs : TYPE
            Description
        activation_fn : TYPE, optional
            Description
        normalizer_fn : TYPE, optional
            Description
        scope : str, optional
            Description
        **kwargs
            Description

        Returns
        -------
        TYPE
            Description
        """
        with tf.variable_scope(scope or 'conv2d_transpose'):
            h = tfl.conv2d_transpose(
                inputs=inputs,
                activation_fn=None,
                normalizer_fn=None,
                weights_initializer=tf.truncated_normal_initializer(stddev=0.02),
                biases_initializer=None,
                **kwargs)
            if normalizer_fn:
                h = normalizer_fn(h)
            if activation_fn:
                h = activation_fn(h)
            return h

    def residual_block(x, n_channels=128, kernel_size=3, scope=None):
        """Summary

        Parameters
        ----------
        x : TYPE
            Description
        n_channels : int, optional
            Description
        kernel_size : int, optional
            Description
        scope : None, optional
            Description

        Returns
        -------
        TYPE
            Description
        """
        with tf.variable_scope(scope or 'residual'):
            h = tf.pad(x, [[0, 0], [1, 1], [1, 1], [0, 0]], "REFLECT")
            h = conv2d(
                inputs=h,
                num_outputs=n_channels,
                kernel_size=kernel_size,
                padding='VALID',
                scope='1')
            h = tf.pad(x, [[0, 0], [1, 1], [1, 1], [0, 0]], "REFLECT")
            h = conv2d(
                inputs=h,
                num_outputs=n_channels,
                kernel_size=kernel_size,
                activation_fn=None,
                padding='VALID',
                scope='2')
            h = tf.add(x, h)
        return h

    def encoder(x, n_filters=32, k_size=3, scope=None):
        """Summary

        Parameters
        ----------
        x : TYPE
            Description
        n_filters : int, optional
            Description
        k_size : int, optional
            Description
        scope : None, optional
            Description

        Returns
        -------
        TYPE
            Description
        """
        with tf.variable_scope(scope or 'encoder'):
            h = tf.pad(x, [[0, 0], [k_size, k_size], [k_size, k_size], [0, 0]],
                       "REFLECT")
            h = conv2d(
                inputs=h,
                num_outputs=n_filters,
                kernel_size=7,
                activation_fn=tf.nn.relu,
                stride=1,
                padding='VALID',
                scope='1')
            h = conv2d(
                inputs=h,
                num_outputs=n_filters * 2,
                kernel_size=k_size,
                activation_fn=tf.nn.relu,
                stride=2,
                scope='2')
            h = conv2d(
                inputs=h,
                num_outputs=n_filters * 4,
                kernel_size=k_size,
                activation_fn=tf.nn.relu,
                stride=2,
                scope='3')
        return h

    def transform(x, img_size=256):
        """Summary

        Parameters
        ----------
        x : TYPE
            Description
        img_size : int, optional
            Description

        Returns
        -------
        TYPE
            Description
        """
        h = x
        if img_size >= 256:
            n_blocks = 9
        else:
            n_blocks = 6
        for block_i in range(n_blocks):
            with tf.variable_scope('block_{}'.format(block_i)):
                h = residual_block(h)
        return h

    def decoder(x, n_filters=32, k_size=3, scope=None):
        """Summary

        Parameters
        ----------
        x : TYPE
            Description
        n_filters : int, optional
            Description
        k_size : int, optional
            Description
        scope : None, optional
            Description

        Returns
        -------
        TYPE
            Description
        """
        with tf.variable_scope(scope or 'decoder'):
            h = conv2d_transpose(
                inputs=x,
                num_outputs=n_filters * 2,
                kernel_size=k_size,
                activation_fn=tf.nn.relu,
                stride=2,
                scope='1')
            h = conv2d_transpose(
                inputs=h,
                num_outputs=n_filters,
                kernel_size=k_size,
                activation_fn=tf.nn.relu,
                stride=2,
                scope='2')
            h = tf.pad(h, [[0, 0], [k_size, k_size], [k_size, k_size], [0, 0]],
                       "REFLECT")
            h = conv2d(
                inputs=h,
                num_outputs=3,
                kernel_size=7,
                stride=1,
                padding='VALID',
                activation_fn=tf.nn.tanh,
                scope='3')
        return h

    def generator(x, scope=None, reuse=None):
        """Summary

        Parameters
        ----------
        x : TYPE
            Description
        scope : None, optional
            Description
        reuse : None, optional
            Description

        Returns
        -------
        TYPE
            Description
        """
        img_size = x.get_shape().as_list()[1]
        with tf.variable_scope(scope or 'generator', reuse=reuse):
            h = encoder(x)
            h = transform(h, img_size)
            h = decoder(h)
        return h

    def discriminator(x, n_filters=64, k_size=4, scope=None, reuse=None):
        """Summary

        Parameters
        ----------
        x : TYPE
            Description
        n_filters : int, optional
            Description
        k_size : int, optional
            Description
        scope : None, optional
            Description
        reuse : None, optional
            Description

        Returns
        -------
        TYPE
            Description
        """
        with tf.variable_scope(scope or 'discriminator', reuse=reuse):
            h = conv2d(
                inputs=x,
                num_outputs=n_filters,
                kernel_size=k_size,
                stride=2,
                normalizer_fn=None,
                scope='1')
            h = conv2d(
                inputs=h,
                num_outputs=n_filters * 2,
                kernel_size=k_size,
                stride=2,
                scope='2')
            h = conv2d(
                inputs=h,
                num_outputs=n_filters * 4,
                kernel_size=k_size,
                stride=2,
                scope='3')
            h = conv2d(
                inputs=h,
                num_outputs=n_filters * 8,
                kernel_size=k_size,
                stride=1,
                scope='4')
            h = conv2d(
                inputs=h,
                num_outputs=1,
                kernel_size=k_size,
                stride=1,
                activation_fn=tf.nn.sigmoid,
                scope='5')
            return h

    def cycle_gan(img_size=256):
        """Summary

        Parameters
        ----------
        img_size : int, optional
            Description

        Returns
        -------
        TYPE
            Description
        """
        X_real = tf.placeholder(
            name='X', shape=[1, img_size, img_size, 3], dtype=tf.float32)
        Y_real = tf.placeholder(
            name='Y', shape=[1, img_size, img_size, 3], dtype=tf.float32)
        X_fake_sample = tf.placeholder(
            name='X_fake_sample',
            shape=[None, img_size, img_size, 3],
            dtype=tf.float32)
        Y_fake_sample = tf.placeholder(
            name='Y_fake_sample',
            shape=[None, img_size, img_size, 3],
            dtype=tf.float32)

        X_fake = generator(Y_real, scope='G_yx')
        Y_fake = generator(X_real, scope='G_xy')
        X_cycle = generator(Y_fake, scope='G_yx', reuse=True)
        Y_cycle = generator(X_fake, scope='G_xy', reuse=True)

        D_X_real = discriminator(X_real, scope='D_X')
        D_Y_real = discriminator(Y_real, scope='D_Y')
        D_X_fake = discriminator(X_fake, scope='D_X', reuse=True)
        D_Y_fake = discriminator(Y_fake, scope='D_Y', reuse=True)
        D_X_fake_sample = discriminator(X_fake_sample, scope='D_X', reuse=True)
        D_Y_fake_sample = discriminator(Y_fake_sample, scope='D_Y', reuse=True)

        # Create losses for generators
        l1 = 10.0
        loss_cycle_X = l1 * l1loss(X_real, X_cycle)
        loss_cycle_Y = l1 * l1loss(Y_real, Y_cycle)
        loss_G_xy = l2loss(D_Y_fake, 1.0)
        loss_G_yx = l2loss(D_X_fake, 1.0)
        loss_G = loss_G_xy + loss_G_yx + loss_cycle_X + loss_cycle_Y

        # Create losses for discriminators
        loss_D_Y = l2loss(D_Y_real, 1.0) + l2loss(D_Y_fake_sample, 0.0)
        loss_D_X = l2loss(D_X_real, 1.0) + l2loss(D_X_fake_sample, 0.0)

        # Summaries for monitoring training
        tf.summary.histogram("D_X_real", D_X_real)
        tf.summary.histogram("D_Y_real", D_Y_real)
        tf.summary.histogram("D_X_fake", D_X_fake)
        tf.summary.histogram("D_Y_fake", D_Y_fake)
        tf.summary.histogram("D_X_fake_sample", D_X_fake_sample)
        tf.summary.histogram("D_Y_fake_sample", D_Y_fake_sample)
        tf.summary.image("X_real", X_real, max_outputs=1)
        tf.summary.image("Y_real", Y_real, max_outputs=1)
        tf.summary.image("X_fake", X_fake, max_outputs=1)
        tf.summary.image("Y_fake", Y_fake, max_outputs=1)
        tf.summary.image("X_cycle", X_cycle, max_outputs=1)
        tf.summary.image("Y_cycle", Y_cycle, max_outputs=1)
        tf.summary.histogram("X_real", X_real)
        tf.summary.histogram("Y_real", Y_real)
        tf.summary.histogram("X_fake", X_fake)
        tf.summary.histogram("Y_fake", Y_fake)
        tf.summary.histogram("X_cycle", X_cycle)
        tf.summary.histogram("Y_cycle", Y_cycle)
        tf.summary.scalar("loss_D_X", loss_D_X)
        tf.summary.scalar("loss_D_Y", loss_D_Y)
        tf.summary.scalar("loss_cycle_X", loss_cycle_X)
        tf.summary.scalar("loss_cycle_Y", loss_cycle_Y)
        tf.summary.scalar("loss_G", loss_G)
        tf.summary.scalar("loss_G_xy", loss_G_xy)
        tf.summary.scalar("loss_G_yx", loss_G_yx)
        summaries = tf.summary.merge_all()

        training_vars = tf.trainable_variables()
        D_X_vars = [v for v in training_vars if v.name.startswith('D_X')]
        D_Y_vars = [v for v in training_vars if v.name.startswith('D_Y')]
        G_xy_vars = [v for v in training_vars if v.name.startswith('G_xy')]
        G_yx_vars = [v for v in training_vars if v.name.startswith('G_yx')]
        G_vars = G_xy_vars + G_yx_vars

        return locals()

    def get_images(path1, path2, img_size=256):
        """Summary

        Parameters
        ----------
        path1 : TYPE
            Description
        path2 : TYPE
            Description
        img_size : int, optional
            Description

        Returns
        -------
        TYPE
            Description
        """
        files1 = [os.path.join(path1, f) for f in os.listdir(path1)]
        files2 = [os.path.join(path2, f) for f in os.listdir(path2)]
        imgs1 = []
        for f in files1:
            try:
                img = imresize(imcrop_tosquare(plt.imread(f)), (img_size, img_size))
            except:
                continue
            if img.ndim == 3:
                imgs1.append(img[..., :3])
            else:
                img = img[..., np.newaxis]
                imgs1.append(np.concatenate([img] * 3, 2))
        imgs1 = np.array(imgs1) / 127.5 - 1.0
        imgs2 = []
        for f in files2:
            try:
                img = imresize(imcrop_tosquare(plt.imread(f)), (img_size, img_size))
            except:
                continue
            if img.ndim == 3:
                imgs2.append(img[..., :3])
            else:
                img = img[..., np.newaxis]
                imgs2.append(np.concatenate([img] * 3, 2))
        imgs2 = np.array(imgs2) / 127.5 - 1.0

        return imgs1, imgs2

    def batch_generator_dataset(imgs1, imgs2):
        """Summary

        Parameters
        ----------
        imgs1 : TYPE
            Description
        imgs2 : TYPE
            Description

        Yields
        ------
        TYPE
            Description
        """
        n_imgs = min(len(imgs1), len(imgs2))
        rand_idxs1 = np.random.permutation(np.arange(len(imgs1)))[:n_imgs]
        rand_idxs2 = np.random.permutation(np.arange(len(imgs2)))[:n_imgs]
        for idx1, idx2 in zip(rand_idxs1, rand_idxs2):
            yield imgs1[[idx1]], imgs2[[idx2]]

    def batch_generator_random_crop(X, Y, min_size=256, max_size=512, n_images=100):
        """Summary

        Parameters
        ----------
        X : TYPE
            Description
        Y : TYPE
            Description
        min_size : int, optional
            Description
        max_size : int, optional
            Description
        n_images : int, optional
            Description

        Yields
        ------
        TYPE
            Description
        """
        r, c, d = X.shape
        Xs, Ys = [], []
        for img_i in range(n_images):
            size = np.random.randint(min_size, max_size)
            max_r = r - size
            max_c = c - size
            this_r = np.random.randint(0, max_r)
            this_c = np.random.randint(0, max_c)
            img = imresize(X[this_r:this_r + size, this_c:this_c + size, :],
                           (min_size, min_size))
            Xs.append(img)
            img = imresize(Y[this_r:this_r + size, this_c:this_c + size, :],
                           (min_size, min_size))
            Ys.append(img)
        imgs1, imgs2 = np.array(Xs) / 127.5 - 1.0, np.array(Ys) / 127.5 - 1.0
        n_imgs = min(len(imgs1), len(imgs2))
        rand_idxs1 = np.random.permutation(np.arange(len(imgs1)))[:n_imgs]
        rand_idxs2 = np.random.permutation(np.arange(len(imgs2)))[:n_imgs]
        for idx1, idx2 in zip(rand_idxs1, rand_idxs2):
            yield imgs1[[idx1]], imgs2[[idx2]]

    def train(ds_X,
              ds_Y,
              ckpt_path='cycle_gan',
              learning_rate=0.0002,
              n_epochs=100,
              img_size=256):
        """Summary

        Parameters
        ----------
        ds_X : TYPE
            Description
        ds_Y : TYPE
            Description
        ckpt_path : str, optional
            Description
        learning_rate : float, optional
            Description
        n_epochs : int, optional
            Description
        img_size : int, optional
            Description
        """
        if ds_X.ndim == 3:
            batch_generator = batch_generator_random_crop
        else:
            batch_generator = batch_generator_dataset

        # How many fake generations to keep around
        capacity = 50

        # Storage for fake generations
        fake_Xs = capacity * [
            np.zeros((1, img_size, img_size, 3), dtype=np.float32)
        ]
        fake_Ys = capacity * [
            np.zeros((1, img_size, img_size, 3), dtype=np.float32)
        ]
        idx = 0
        it_i = 0

        # Train
        with tf.Graph().as_default(), tf.Session() as sess:
            # Load the network
            net = cycle_gan(img_size=img_size)

            # Build optimizers
            D_X = tf.train.AdamOptimizer(learning_rate=learning_rate).minimize(
                net['loss_D_X'], var_list=net['D_X_vars'])
            D_Y = tf.train.AdamOptimizer(learning_rate=learning_rate).minimize(
                net['loss_D_Y'], var_list=net['D_Y_vars'])
            G = tf.train.AdamOptimizer(learning_rate=learning_rate).minimize(
                net['loss_G'], var_list=net['G_vars'])

            init_op = tf.group(tf.global_variables_initializer(),
                               tf.local_variables_initializer())
            sess.run(init_op)
            saver = tf.train.Saver()
            writer = tf.summary.FileWriter(ckpt_path)
            for epoch_i in range(n_epochs):
                for X, Y in batch_generator(ds_X, ds_Y):

                    # First generate in both directions
                    X_fake, Y_fake = sess.run(
                        [net['X_fake'], net['Y_fake']],
                        feed_dict={net['X_real']: X,
                                   net['Y_real']: Y})

                    # Now sample from history
                    if it_i < capacity:
                        # Not enough samples yet, fill up history buffer
                        fake_Xs[idx] = X_fake
                        fake_Ys[idx] = Y_fake
                        idx = (idx + 1) % capacity
                    elif np.random.random() > 0.5:
                        # Swap out a random idx from history
                        rand_idx = np.random.randint(0, capacity)
                        fake_Xs[rand_idx], X_fake = X_fake, fake_Xs[rand_idx]
                        fake_Ys[rand_idx], Y_fake = Y_fake, fake_Ys[rand_idx]
                    else:
                        # Use current generation
                        pass

                    # Optimize G Networks
                    loss_G = sess.run(
                        [net['loss_G'], G],
                        feed_dict={
                            net['X_real']: X,
                            net['Y_real']: Y,
                            net['Y_fake_sample']: Y_fake,
                            net['X_fake_sample']: X_fake
                        })[0]

                    # Optimize D_Y
                    loss_D_Y = sess.run(
                        [net['loss_D_Y'], D_Y],
                        feed_dict={
                            net['X_real']: X,
                            net['Y_real']: Y,
                            net['Y_fake_sample']: Y_fake
                        })[0]

                    # Optimize D_X
                    loss_D_X = sess.run(
                        [net['loss_D_X'], D_X],
                        feed_dict={
                            net['X_real']: X,
                            net['Y_real']: Y,
                            net['X_fake_sample']: X_fake
                        })[0]

                    print(it_i, 'G:', loss_G, 'D_X:', loss_D_X, 'D_Y:', loss_D_Y)

                    # Update summaries
                    if it_i % 100 == 0:
                        summary = sess.run(
                            net['summaries'],
                            feed_dict={
                                net['X_real']: X,
                                net['Y_real']: Y,
                                net['X_fake_sample']: X_fake,
                                net['Y_fake_sample']: Y_fake
                            })
                        writer.add_summary(summary, it_i)
                    it_i += 1

                # Save
                if epoch_i % 50 == 0:
                    saver.save(
                        sess,
                        os.path.join(ckpt_path, 'model.ckpt'),
                        global_step=epoch_i)