Add files via upload

Author: ZZUTK

Adapt_CAAE/ResLearn.py

@@ -0,0 +1,78 @@
+# check GPU status
+from pynvml import *
+from time import sleep
+from datetime import datetime
+
+
+def gpu_memory_usage(is_print=True):
+    try:
+        nvmlInit()
+        # version = nvmlSystemGetDriverVersion()
+        deviceCount = nvmlDeviceGetCount()
+        GPU = {}
+        for i in range(deviceCount):
+            GPU[i] = {}
+            handle = nvmlDeviceGetHandleByIndex(i)
+            info = nvmlDeviceGetMemoryInfo(handle)
+            GPU[i]['total'] = info.total / 1024.0 / 1024.0 / 1024.0
+            GPU[i]['free'] = info.free / 1024.0 / 1024.0 / 1024.0
+            if is_print:
+                print("\nGPU #%d Memory Usage:"
+                      "\n\tTotal:\t%4.2fGB\n\tFree:\t%4.2fGB" %
+                      (i, GPU[i]['total'], GPU[i]['free']))
+                print datetime.now()
+        nvmlShutdown()
+        return GPU
+    except:
+        print "Fail to check GPU status!"
+        exit(0)
+
+
+def auto_queue(gpu_memory_require=3.2, interval=1, schedule=None):
+    # input arg: schedule = datetime(year, month, day, hour, minute, second)
+    if schedule is None:
+        schedule = datetime.today()
+    else:
+        print '\nScheduled time: ', schedule
+
+    # wait until the scheduled time
+    now = datetime.today()
+    while now.year < schedule.year or now.month < schedule.month or now.day < schedule.day or \
+          now.hour < schedule.hour or now.minute < schedule.minute or now.second < schedule.second:
+        now = datetime.today()
+        sleep(interval)
+
+    gpu_stat = gpu_memory_usage()
+    if gpu_stat[0]['total'] < gpu_memory_require:
+        print 'Memory requirement is larger than GPU total memory'
+        exit(1)
+    while gpu_stat[0]['free'] < gpu_memory_require:
+        sleep(interval)  # second
+        gpu_stat = gpu_memory_usage()
+    return gpu_stat
+
+
+def set_memory_usage(usage=12.0, allow_growth=True):
+    auto_queue(gpu_memory_require=usage)
+    try:
+        import tensorflow as tf
+        assert type(usage) is int or float
+        assert usage >= 0
+       
+        config = tf.ConfigProto()
+        gpu_stat = gpu_memory_usage()
+        total_memory = gpu_stat[0]['total']
+        if usage > total_memory:
+            usage_percentage = 1.0
+        else:
+            usage_percentage = usage / total_memory
+        config.gpu_options.allow_growth = allow_growth
+        config.gpu_options.per_process_gpu_memory_fraction = usage_percentage
+        return config
+    except:
+        print 'Failed to set memory usage!'
+        return None
+
+
+if __name__ == '__main__':
+    gpu_memory_usage()

Adapt_CAAE/checkGPU.py

@@ -0,0 +1,76 @@
+import tensorflow as tf
+from ResLearn import ResLearn
+import checkGPU
+import numpy as np
+
+
+flags = tf.app.flags
+flags.DEFINE_integer(flag_name='epoch', default_value=50, docstring='number of epochs')
+flags.DEFINE_integer(flag_name='batch_size', default_value=25, docstring='batch size')
+flags.DEFINE_integer(flag_name='is_train', default_value=1, docstring='training mode')
+flags.DEFINE_integer(flag_name='is_bn', default_value=0, docstring='enable batch normalization')
+flags.DEFINE_string(flag_name='dataset', default_value='UTKFace', docstring='dataset name')
+flags.DEFINE_string(flag_name='savedir', default_value='save', docstring='dir for saving training results')
+flags.DEFINE_string(flag_name='testdir', default_value='None', docstring='dir for testing images')
+flags.DEFINE_float(flag_name='param0', default_value=1.0, docstring='weight of discriminator loss on fake images')
+flags.DEFINE_float(flag_name='param1', default_value=1.0, docstring='weight of reconstruct loss on artifact modeling')
+flags.DEFINE_integer(flag_name='is_schedule', default_value=0, docstring='scheduled running')
+flags.DEFINE_integer(flag_name='day', default_value=1, docstring='date')
+flags.DEFINE_integer(flag_name='hr', default_value=0, docstring='hour')
+flags.DEFINE_integer(flag_name='min', default_value=0, docstring='minute')
+
+FLAGS = flags.FLAGS
+
+
+gpu_memory_require = 7.0
+
+
+def main(_):
+    from datetime import datetime
+    if FLAGS.is_schedule:
+        today = datetime.today()
+        checkGPU.auto_queue(
+            gpu_memory_require=gpu_memory_require,
+            interval=1,
+            schedule=datetime(year=today.year, month=today.month, day=FLAGS.day, hour=FLAGS.hr, minute=FLAGS.min)
+        )
+    config = checkGPU.set_memory_usage(
+        usage=gpu_memory_require,
+        allow_growth=True
+    )
+
+    # print settings
+    import pprint
+    pprint.pprint(FLAGS.__flags)
+
+    with tf.Session(config=config) as session:
+        model = ResLearn(
+            session,  # TensorFlow session
+            is_training=FLAGS.is_train,  # flag for training or testing mode
+            save_dir=FLAGS.savedir,  # path to save checkpoints, samples, and summary
+            dataset_name=FLAGS.dataset,  # name of the dataset in the folder ./data
+            size_batch=FLAGS.batch_size,
+            enable_bn=FLAGS.is_bn
+        )
+        if FLAGS.is_train:
+            print '\n\tTraining Mode'
+            model.train(
+                num_epochs=FLAGS.epoch,  # number of epochs
+                params=[FLAGS.param0, FLAGS.param1]
+            )
+        else:
+            print '\n\tTesting Mode'
+            model.custom_test(
+                testing_samples_dir=FLAGS.testdir + '/*jpg'
+            )
+
+
+if __name__ == '__main__':
+    if 0:
+        print 'Run on CPU'
+        with tf.device("/cpu:0"):
+            gpu_memory_require = 0.0
+            tf.app.run()
+
+    tf.app.run()
+

Adapt_CAAE/data/UTKFace/100_1_2_20170112213615815.jpg.chip.jpg

@@ -0,0 +1,137 @@
+from __future__ import division
+import tensorflow as tf
+import numpy as np
+from scipy.misc import imread, imresize, imsave
+
+
+class batch_norm(object):
+    def __init__(self, epsilon=1e-5, momentum=0.9, name="batch_norm"):
+        with tf.variable_scope(name):
+            self.epsilon = epsilon
+            self.momentum = momentum
+            self.name = name
+
+    def __call__(self, x, train=True):
+        return tf.contrib.layers.batch_norm(x, decay=self.momentum, updates_collections=None, epsilon=self.epsilon,
+                                            scale=True, scope=self.name)
+
+
+def conv2d(input_map, num_output_channels, size_kernel=5, stride=2, name='conv2d'):
+    with tf.variable_scope(name):
+        stddev = np.sqrt(2.0 / (np.sqrt(input_map.get_shape()[-1].value * num_output_channels) * size_kernel ** 2))
+        stddev = 0.02
+        kernel = tf.get_variable(
+            name='w',
+            shape=[size_kernel, size_kernel, input_map.get_shape()[-1], num_output_channels],
+            dtype=tf.float32,
+            initializer=tf.truncated_normal_initializer(stddev=stddev)
+        )
+        biases = tf.get_variable(
+            name='b',
+            shape=[num_output_channels],
+            dtype=tf.float32,
+            initializer=tf.constant_initializer(0.0)
+        )
+        conv = tf.nn.conv2d(input_map, kernel, strides=[1, stride, stride, 1], padding='SAME')
+        return tf.nn.bias_add(conv, biases)
+
+
+def fc(input_vector, num_output_length, name='fc'):
+    with tf.variable_scope(name):
+        stddev = np.sqrt(1.0 / (np.sqrt(input_vector.get_shape()[-1].value * num_output_length)))
+        stddev = 0.02
+        w = tf.get_variable(
+            name='w',
+            shape=[input_vector.get_shape()[1], num_output_length],
+            dtype=tf.float32,
+            initializer=tf.random_normal_initializer(stddev=stddev)
+        )
+        b = tf.get_variable(
+            name='b',
+            shape=[num_output_length],
+            dtype=tf.float32,
+            initializer=tf.constant_initializer(0.0)
+        )
+        return tf.matmul(input_vector, w) + b
+
+
+def deconv2d(input_map, output_shape, size_kernel=5, stride=2, stddev=0.02, name='deconv2d'):
+    with tf.variable_scope(name):
+        stddev = np.sqrt(1.0 / (np.sqrt(input_map.get_shape()[-1].value * output_shape[-1]) * size_kernel ** 2))
+        stddev = 0.02
+        # filter : [height, width, output_channels, in_channels]
+        kernel = tf.get_variable(
+            name='w',
+            shape=[size_kernel, size_kernel, output_shape[-1], input_map.get_shape()[-1]],
+            dtype=tf.float32,
+            initializer=tf.random_normal_initializer(stddev=stddev)
+        )
+        biases = tf.get_variable(
+            name='b',
+            shape=[output_shape[-1]],
+            dtype=tf.float32,
+            initializer=tf.constant_initializer(0.0)
+        )
+        deconv = tf.nn.conv2d_transpose(input_map, kernel, strides=[1, stride, stride, 1], output_shape=output_shape)
+        return tf.nn.bias_add(deconv, biases)
+       
+
+def lrelu(logits, leak=0.2):
+    return tf.maximum(logits, leak*logits)
+
+
+def concat_label(x, label, duplicate=1):
+    x_shape = x.get_shape().as_list()
+    if duplicate < 1:
+        return x
+    # duplicate the label to enhance its effect, does it really affect the result?
+    label = tf.tile(label, [1, duplicate])
+    label_shape = label.get_shape().as_list()
+    if len(x_shape) == 2:
+        return tf.concat(1, [x, label])
+    elif len(x_shape) == 4:
+        label = tf.reshape(label, [x_shape[0], 1, 1, label_shape[-1]])
+        return tf.concat(3, [x, label*tf.ones([x_shape[0], x_shape[1], x_shape[2], label_shape[-1]])])
+
+
+def load_image(
+        image_path,  # path of a image
+        image_size=64,  # expected size of the image
+        image_value_range=(-1, 1),  # expected pixel value range of the image
+        is_gray=False,  # gray scale or color image
+):
+    if is_gray:
+        image = imread(image_path, flatten=True).astype(np.float32)
+    else:
+        image = imread(image_path).astype(np.float32)
+    image = imresize(image, [image_size, image_size])
+    image = image.astype(np.float32) * (image_value_range[-1] - image_value_range[0]) / 255.0 + image_value_range[0]
+    return image
+
+
+def save_batch_images(
+        batch_images,   # a batch of images
+        save_path,  # path to save the images
+        image_value_range=(-1,1),   # value range of the input batch images
+        size_frame=None     # size of the image matrix, number of images in each row and column
+):
+    # transform the pixcel value to 0~1
+    images = (batch_images - image_value_range[0]) / (image_value_range[-1] - image_value_range[0])
+    if size_frame is None:
+        auto_size = int(np.ceil(np.sqrt(images.shape[0])))
+        size_frame = [auto_size, auto_size]
+    img_h, img_w = batch_images.shape[1], batch_images.shape[2]
+    frame = np.zeros([img_h * size_frame[0], img_w * size_frame[1], 3])
+    for ind, image in enumerate(images):
+        ind_col = ind % size_frame[1]
+        ind_row = ind // size_frame[1]
+        frame[(ind_row * img_h):(ind_row * img_h + img_h), (ind_col * img_w):(ind_col * img_w + img_w), :] = image
+    imsave(save_path, np.clip(frame, 0.0, 1.0))  # imsave(save_path, frame)
+
+
+
+
+
+
+
+