som.py

from random import choice as random_choice
from argparse import ArgumentParser

import numpy as np
import tensorflow as tf

import cv2

from clustering.som import SelfOrganizingMap


def main():
    parser = ArgumentParser()
    parser.add_argument('file', metavar='path', help='The TFRecord file to load.')
    args = parser.parse_args()

    random_seed = None

    width = 128
    height = 128
    dimensions = 32*32*3
    total_iterations = 1000

    np.random.seed(random_seed)
    with tf.Graph().as_default() as graph:

        with tf.variable_scope('input'):
            # TODO: glob the actual TFRecord files
            queue = tf.train.string_input_producer([args.file])

            options = tf.python_io.TFRecordOptions(tf.python_io.TFRecordCompressionType.GZIP)
            reader = tf.TFRecordReader(options=options)
            _, example = reader.read(queue)
            features = tf.parse_single_example(example, features={'raw': tf.FixedLenFeature([], tf.string)})
            image = tf.decode_raw(features['raw'], tf.uint8)
            raw_image = tf.reshape(image, (32, 32, 3))

        with tf.variable_scope('som'):
            som = SelfOrganizingMap(height, width, dimensions, total_iterations, random_seed=random_seed)

    coord = tf.train.Coordinator()
    with tf.Session(graph=graph) as sess:
        # initialize all variables
        init = tf.group(tf.global_variables_initializer(), tf.local_variables_initializer())
        sess.run(init)

        # fire up the input queue runners
        threads = tf.train.start_queue_runners(coord=coord)

        # work with real data ...
        data = [(1., 0., 0.),
                (0., 1., 0.),
                (0., 0., 1.),
                (1., 1., 0.),
                (0., 1., 1.),
                (1., 0., 1.),
                (1., 1., 1.),
                (0., 0., 0.),
                (.5, .5, .5)]

        # ... or work with synthetic data
        #data = np.random.rand(512, 3)

        window_title = 'Grid'
        cv2.namedWindow(window_title, cv2.WINDOW_NORMAL)

        # prepare the output buffer
        y_stride = 2
        x_stride = 2
        image = np.empty((32 * height // y_stride, 32 * width // x_stride, 3))

        canceled = False
        while not coord.should_stop():
            for _ in range(0, 100):
                if not coord.should_stop():
                    # fetch one example
                    target = sess.run(raw_image)
                    target = np.reshape(target, newshape=(dimensions,)) / 255.

                    # run one iteration
                    iteration = som.run(target, sess)
                    if iteration >= total_iterations:
                        coord.should_stop()
                        continue

            # reporting every couple of iterations
            if iteration % 10 == 0:
                print('Displaying iteration %i ...' % iteration)
                grid = som.get_grid(sess)

                for y in range(0, height, y_stride):
                    yy = np.arange(y*32//y_stride, y*32//y_stride+32)
                    for x in range(0, width, x_stride):
                        grid_image = np.reshape(grid[y, x], (32, 32, 3))
                        xx = x*32//x_stride
                        image[yy, xx:xx+32, :] = grid_image

                cv2.imshow(window_title, image)
                if cv2.waitKey(10) & 0xff == 27:
                    canceled = True
                    break

        coord.request_stop()
        coord.join(threads)

        if not canceled:
            print('Done.')
            cv2.waitKey(0)
        cv2.destroyAllWindows()

    pass


if __name__ == '__main__':
    main()