cae.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]

            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)

                        z = tf.nn.conv2d(next_layer_input, weight, strides=(1, 1, 1, 1), padding='SAME') + bias
                        z = tf.nn.max_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

            # 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):
                        layers_in = int(next_layer_input.get_shape()[3])
                        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)

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

                        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)

                        z = 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'):
            train_step = tf.train.AdamOptimizer(0.01).minimize(cae.loss, global_step=global_step)
            tf.summary.scalar('loss', cae.loss)

        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')
    with sv.managed_session() as sess:

        start_time = time()

        # 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 = sess.run([train_step, cae.loss, global_step])

            # GUI event loop handling
            cv2.waitKey(1)

            if time() - start_time >= 2.:
                r, summary, loss = sess.run((reconstruction, image_summaries, cae.loss), feed_dict={x: t_sample})
                print('%5i: loss %f' % (i, 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()