cae_tw.py

# https://gist.github.com/saliksyed/593c950ba1a3b9dd08d5

from time import time
from glob import glob
from collections import namedtuple

import tensorflow as tf
import numpy as np
import cv2

from clustering.import_stage import build_import_stage
from clustering.autoencoder import AutoencoderOutput


def deep_test():
    def sample_image():
        sample = np.ones((32, 32, 3), dtype=float)
        sample = cv2.line(sample, (0, 0), (32, 32), (0., 0., 1.), lineType=cv2.LINE_AA)
        sample = cv2.line(sample, (0, 32), (32, 0), (0., 1., 0.), lineType=cv2.LINE_AA)
        sample = cv2.circle(sample, (16, 16), 14, (1., 0., 0.), lineType=cv2.LINE_AA)
        return sample

    def sample_window(input, output, timeout=0, window=None):
        if input is None:
            input = sample_image()

        if output is None:
            output = input

        canvas = np.zeros((40, 80, 3))
        canvas[4:36, 4:36, :] = input
        canvas[4:36, 44:76, :] = output
        canvas = cv2.resize(canvas, (320, 160), interpolation=cv2.INTER_NEAREST)

        # render
        name = window
        if window is window:
            name = 'image'
            cv2.namedWindow(name)

        cv2.imshow(name, canvas)
        cv2.waitKey(timeout)

        if window is None:
            cv2.destroyAllWindows()

    with tf.Graph().as_default() as graph:
        files = glob('data/*.tfrecord.gz')
        input_batch = build_import_stage(files, num_epochs=None)

        x = input_batch

        Definition = namedtuple('Definition', ['kernel_shape', 'layers_out', 'pooling'])

        input_def = Definition(kernel_shape=(5, 5), layers_out=3, pooling=1)
        layer_sizes = [Definition(kernel_shape=(3, 3), layers_out=64, pooling=2),      # 32x32 -> 16x16
                       Definition(kernel_shape=(3, 3), layers_out=128, pooling=2),     # 16x16 -> 8x8
                       Definition(kernel_shape=(3, 3), layers_out=64, pooling=2)]      # 8x8 -> 4x4
                       #Definition(kernel_shape=(3, 3), layers_out=1024, pooling=2),   # 4x4 -> 2x2
                       #Definition(kernel_shape=(3, 3), layers_out=512,  pooling=2)]   # 2x2 -> 1x1

        with tf.variable_scope('autoencoder'):
            next_layer_input = x
            batch_size = tf.shape(next_layer_input)[0]

            all_weights = []

            with tf.variable_scope('encoder'):
                for i, dim in enumerate(layer_sizes):
                    with tf.variable_scope('encode_%i' % i):
                        layers_in = int(next_layer_input.get_shape()[3])

                        height, width = dim.kernel_shape
                        kernel_shape = (height, width, layers_in, dim.layers_out)

                        initial_weight = tf.truncated_normal_initializer()
                        initial_bias = tf.constant(0., dtype=tf.float32, shape=(dim.layers_out,))

                        weight = tf.get_variable(name='W', shape=kernel_shape, initializer=initial_weight)
                        bias = tf.get_variable(name='b', initializer=initial_bias)

                        all_weights.append(weight)

                        z = tf.nn.bias_add(tf.nn.conv2d(next_layer_input, weight,
                                                        strides=(1, dim.pooling, dim.pooling, 1), padding='SAME'),
                                           bias)

                        # z = tf.nn.bias_add(tf.nn.conv2d(next_layer_input, weight, strides=(1, 1, 1, 1), padding='SAME'),
                        #                    bias)
                        # z = tf.nn.avg_pool(z, ksize=(1, dim.pooling, dim.pooling, 1),
                        #                    strides=(1, dim.pooling, dim.pooling, 1), padding='SAME')
                        h = tf.nn.sigmoid(z)

                    next_layer_input = h

            # the low-dimension encoded value
            embedding = tf.identity(next_layer_input, name='embedding')
            next_layer_input = embedding

            # reverse the weights for the decoders
            all_weights.reverse()

            # build the reconstruction layers by reversing the reductions
            layer_sizes.reverse()
            layer_sizes.append(input_def)

            with tf.variable_scope('decoder'):
                for i, (dim, next) in enumerate(zip(layer_sizes[1:], layer_sizes[0:-1])):
                    with tf.variable_scope('decode_%i' % i):
                        height_in = int(next_layer_input.get_shape()[1])
                        width_in = int(next_layer_input.get_shape()[2])
                        output_shape = (batch_size, next.pooling * height_in, next.pooling * width_in, dim.layers_out)

                        initial_bias = tf.constant(0., dtype=tf.float32, shape=(dim.layers_out,))

                        weight = all_weights[i]
                        bias = tf.get_variable(name='b', initializer=initial_bias)

                        z = tf.nn.bias_add(tf.nn.conv2d_transpose(next_layer_input, weight, output_shape=output_shape,
                                           strides=(1, next.pooling, next.pooling, 1)), bias)
                        h = tf.nn.sigmoid(z)
                        h = tf.reshape(h, output_shape)

                    next_layer_input = h

            reconstruction = tf.identity(next_layer_input, name='reconstruction')

            with tf.variable_scope('optimization'):
                loss = tf.reduce_mean(tf.square(x - reconstruction), name='loss')

        cae = AutoencoderOutput(embedding=embedding, reconstruction=reconstruction, loss=loss)

        global_step = tf.Variable(initial_value=0, trainable=False, name='global_step')
        with tf.name_scope('training'):
            t_learning_rate = tf.Variable(initial_value=tf.constant(1e-3, dtype=tf.float32), dtype=tf.float32, name='learning_rate')
            step_down_lr = tf.assign(t_learning_rate, t_learning_rate * 0.1, name='lr_step_down')

            train_step = tf.train.AdamOptimizer(t_learning_rate).minimize(cae.loss, global_step=global_step)

            tf.summary.scalar('loss', cae.loss)
            tf.summary.scalar('learning_rate', t_learning_rate)

        tf.summary.image('input', x, collections=['snapshot'])
        tf.summary.image('reconstruction', cae.reconstruction, collections=['snapshot'])
        image_summaries = tf.summary.merge_all('snapshot')

    sv = tf.train.Supervisor(graph=graph, logdir='log', save_summaries_secs=20)
    with sv.managed_session() as sess:

        start_time = time()

        # learning rate schedule
        step_down_every = 1000
        next_stop = step_down_every

        # initial metric
        loss, i = sess.run([cae.loss, global_step])
        print('%5i: loss %f' % (i, loss))

        # run a sample image
        sample = cv2.imread('sample/patch.jpg')
        t_sample = cv2.cvtColor(sample, cv2.COLOR_BGR2RGB).astype(float) / 255.
        t_sample = np.expand_dims(t_sample, 0)
        sample = sample.astype(float) / 255.

        r = sess.run(reconstruction, feed_dict={x: t_sample})

        window = 'image'
        sample_window(sample, r[0, :, :, :], timeout=1, window=window)

        while not sv.should_stop():
            _, loss, i, summary = sess.run([train_step, cae.loss, global_step, image_summaries])
            sv.summary_computed(sess, summary, global_step=i)

            if i > next_stop:
                next_stop += step_down_every
                learning_rate = sess.run(step_down_lr)
                print('Stepping down the learning rate to %f.' % learning_rate)

            # GUI event loop handling
            cv2.waitKey(1)

            if time() - start_time >= 2.:
                r, summary, v_loss = sess.run((reconstruction, image_summaries, cae.loss),
                                              feed_dict={x: t_sample})
                print('%5i: training loss: %f, validation loss %f' % (i, loss, v_loss))

                sv.summary_computed(sess, summary, global_step=i)

                r = np.clip(r[0, :, :, :] * 255., 0., 255.).astype(np.uint8)
                r = cv2.cvtColor(r, cv2.COLOR_RGB2BGR).astype(float) / 255.
                sample_window(sample, r, timeout=1, window=window)

                start_time = time()

        cv2.destroyAllWindows()


if __name__ == '__main__':
    # simple_test()
    deep_test()