├── SRM_Kernels.npy
├── README.md
├── main.py
├── queues.py
├── YeNet.py
├── generator.py
├── utils.py
└── layers.py


/SRM_Kernels.npy:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/Caenorst/YeNet-Tensorflow/HEAD/SRM_Kernels.npy


--------------------------------------------------------------------------------
/README.md:
--------------------------------------------------------------------------------
1 | # YeNet-Tensorflow
2 | Tensorflow implementation of "Deep Learning Hierarchical Representation for Image Steganalysis" by Jian Ye, Jiangqun Ni and Yang Yi
3 | 
4 | ## Dataset
5 | Training and validation contains there own cover and stego / beta maps images directory, stego images must have the same name than there corresponding cover
6 | 
7 | ## Publication
8 | [The publication can be found here](http://ieeexplore.ieee.org/document/7937836/)
9 | 


--------------------------------------------------------------------------------
/main.py:
--------------------------------------------------------------------------------
 1 | import argparse
 2 | import numpy as np
 3 | import tensorflow as tf
 4 | from functools import partial
 5 | 
 6 | from utils import *
 7 | from generator import *
 8 | from queues import *
 9 | from YeNet import YeNet
10 | 
11 | parser = argparse.ArgumentParser(description='PyTorch implementation of YeNet')
12 | parser.add_argument('train_cover_dir', type=str, metavar='PATH',
13 |                     help='path of directory containing all ' +
14 |                     'training cover images')
15 | parser.add_argument('train_stego_dir', type=str, metavar='PATH',
16 |                     help='path of directory containing all ' +
17 |                     'training stego images or beta maps')
18 | parser.add_argument('valid_cover_dir', type=str, metavar='PATH',
19 |                     help='path of directory containing all ' +
20 |                     'validation cover images')
21 | parser.add_argument('valid_stego_dir', type=str, metavar='PATH',
22 |                     help='path of directory containing all ' +
23 |                     'validation stego images or beta maps')
24 | parser.add_argument('--batch-size', type=int, default=32, metavar='N',
25 |                     help='input batch size for training (default: 32)')
26 | parser.add_argument('--test-batch-size', type=int, default=32, metavar='N',
27 |                     help='input batch size for testing (default: 32)')
28 | parser.add_argument('--epochs', type=int, default=1000, metavar='N',
29 |                     help='number of epochs to train (default: 1000)')
30 | parser.add_argument('--lr', type=float, default=4e-1, metavar='LR',
31 |                     help='learning rate (default: 4e-1)')
32 | parser.add_argument('--use-batch-norm', action='store_true', default=False,
33 |                     help='use batch normalization after each activation,' +
34 |                     ' also disable pair constraint (default: False)')
35 | parser.add_argument('--embed-otf', action='store_true', default=False,
36 |                     help='use beta maps and embed on the fly instead' +
37 |                     ' of use stego images (default: False)')
38 | parser.add_argument('--no-cuda', action='store_true', default=False,
39 |                     help='disables CUDA training')
40 | parser.add_argument('--gpu', type=int, default=0,
41 |                     help='index of gpu used (default: 0)')
42 | parser.add_argument('--seed', type=int, default=1, metavar='S',
43 |                     help='random seed (default: 1)')
44 | parser.add_argument('--log-interval', type=int, default=200, metavar='N',
45 |                     help='how many batches to wait ' +
46 |                     'before logging training status')
47 | parser.add_argument('--log-path', type=str, default='logs/',
48 |                     metavar='PATH', help='path to generated log file')
49 | args = parser.parse_args()
50 | 
51 | import os
52 | os.environ['CUDA_DEVICE_ORDER'] = 'PCI_BUS_ID'
53 | os.environ['CUDA_VISIBLE_DEVICES'] = '' if args.no_cuda else str(args.gpu)
54 | 
55 | tf.set_random_seed(args.seed)
56 | train_ds_size = len(glob(args.train_cover_dir + '/*')) * 2
57 | if args.embed_otf:
58 |     train_gen = partial(gen_embedding_otf, args.train_cover_dir, \
59 |                         args.train_stego_dir, args.use_batch_norm)
60 | else:
61 |     train_gen = partial(gen_flip_and_rot, args.train_cover_dir, \
62 |                         args.train_stego_dir,  args.use_batch_norm)
63 | 
64 | valid_ds_size = len(glob(args.valid_cover_dir + '/*')) * 2
65 | valid_gen = partial(gen_valid, args.valid_cover_dir, \
66 |                     args.valid_stego_dir)
67 | 
68 | if valid_ds_size % 32 != 0:
69 |     raise ValueError("change batch size for validation")
70 |     
71 | optimizer = tf.train.AdadeltaOptimizer(args.lr)
72 | 
73 | 
74 | train(YeNet, train_gen, valid_gen, args.batch_size, \
75 |       args.test_batch_size, valid_ds_size, \
76 |       optimizer, args.log_interval, train_ds_size, \
77 |       args.epochs * train_ds_size, train_ds_size, args.log_path, 8)
78 | 


--------------------------------------------------------------------------------
/queues.py:
--------------------------------------------------------------------------------
  1 | import numpy as np
  2 | import tensorflow as tf
  3 | import threading
  4 | import h5py
  5 | import functools
  6 | 
  7 | def hdf5baseGen(filepath, thread_idx, n_threads):
  8 |     with h5py.File(filepath, 'r') as f:
  9 |         keys = f.keys()
 10 |         nb_data = f[keys[0]].shape[0]
 11 |         idx = thread_idx
 12 |         while True:
 13 |             yield [np.expand_dims(f[key][idx], 0) for key in keys]
 14 |             idx = (idx + n_threads) % nb_data
 15 | 
 16 | class GeneratorRunner():
 17 |     """
 18 |     This class manage a multithreaded queue filled with a generator
 19 |     """
 20 |     def __init__(self, generator, capacity):
 21 |         """
 22 |         inputs: generator feeding the data, must have thread_idx 
 23 |             as parameter (but the parameter may be not used)
 24 |         """
 25 |         self.generator = generator
 26 |         _input = generator(0,1).next()
 27 |         if type(_input) is not list:
 28 |             raise ValueError("generator doesn't return" \
 29 |                              "a list: %r" %  type(_input))
 30 |         input_batch_size = _input[0].shape[0]
 31 |         if not all(_input[i].shape[0] == input_batch_size for i in range(len(_input))):
 32 |             raise ValueError("all the inputs doesn't have " + \
 33 |             				 "the same batch size," \
 34 |                              "the batch sizes are: %s" % [_input[i].shape[0] \
 35 |                              for i in range(len(_input))])
 36 |         self.data = []
 37 |         self.dtypes = []
 38 |         self.shapes = []
 39 |         for i in range(len(_input)):
 40 |             self.shapes.append(_input[i].shape[1:])
 41 |             self.dtypes.append(_input[i].dtype)
 42 |             self.data.append(tf.placeholder(dtype=self.dtypes[i], \
 43 |                                 shape=(input_batch_size,) + self.shapes[i]))
 44 |         self.queue = tf.FIFOQueue(capacity, shapes=self.shapes, \
 45 |                                 dtypes=self.dtypes)
 46 |         self.enqueue_op = self.queue.enqueue_many(self.data)
 47 |         self.close_queue_op = self.queue.close(cancel_pending_enqueues=True)
 48 |         
 49 |     def get_batched_inputs(self, batch_size):
 50 |         """
 51 |         Return tensors containing a batch of generated data
 52 |         """
 53 |         batch = self.queue.dequeue_many(batch_size)
 54 |         return batch
 55 |     
 56 |     def thread_main(self, sess, thread_idx=0, n_threads=1):
 57 |         try:
 58 |             for data in self.generator(thread_idx, n_threads):
 59 |                 sess.run(self.enqueue_op, feed_dict={i: d \
 60 |                             for i, d in zip(self.data, data)})
 61 |                 if self.stop_threads:
 62 |                     return
 63 |         except RuntimeError:
 64 |             pass
 65 |         except tf.errors.CancelledError:
 66 |             pass
 67 |     
 68 |     def start_threads(self, sess, n_threads=1):
 69 |         self.stop_threads = False
 70 |         self.threads = []
 71 |         for n in range(n_threads):
 72 |             t = threading.Thread(target=self.thread_main, args=(sess, n, n_threads))
 73 |             t.daemon = True
 74 |             t.start()
 75 |             self.threads.append(t)
 76 |         return self.threads
 77 | 
 78 |     def stop_runner(self, sess):
 79 |         self.stop_threads = True
 80 |         # j = 0
 81 |         # while np.any([t.is_alive() for t in self.threads]):
 82 |         #     j += 1
 83 |         #     if j % 100 = 0:
 84 |         #         print [t.is_alive() for t in self.threads]
 85 |         sess.run(self.close_queue_op)
 86 | 
 87 | def queueSelection(runners, sel, batch_size):
 88 |     selection_queue = tf.FIFOQueue.from_list(sel, [r.queue for r in runners])
 89 |     return selection_queue.dequeue_many(batch_size)
 90 | 
 91 | def doubleQueue(runner1, runner2, is_runner1, batch_size1, batch_size2):
 92 |     return tf.cond(is_runner1, lambda: runner1.queue.dequeue_many(batch_size1), \
 93 |                 lambda: runner2.queue.dequeue_many(batch_size2))
 94 | 
 95 | if __name__ == '__main__':
 96 |     def randomGen(img_size, enqueue_batch_size, thread_idx, n_threads):
 97 |         while True:
 98 |             batch_of_1_channel_imgs = np.random.rand(enqueue_batch_size, \
 99 |                                                      img_size, img_size, 1)
100 |             batch_of_labels = np.random.randint(0,11,enqueue_batch_size)
101 |             return [batch_of_1_channel_imgs, batch_of_labels]
102 | 
103 |     TRAIN_BATCH_SIZE = 64
104 |     VALID_BATCH_SIZE = 10
105 |     train_runner = GeneratorRunner(functool.partial(randomGen, \
106 |                                    (128, 10)), TRAIN_BATCH_SIZE * 10)
107 |     valid_runner = GeneratorRunner(functool.partial(randomGen, \
108 |                                    (128, 10)), VALID_BATCH_SIZE * 10)
109 |     is_training = tf.Variable(True)
110 |     batch_size = tf.Variable(TRAIN_BATCH_SIZE)
111 |     enable_training_op = tf.group(tf.assign(is_training, True), \
112 |                                   tf.assign(batch_size, TRAIN_BATCH_SIZE))
113 |     disable_training_op = tf.group(tf.assign(is_training, False), \
114 |                                   tf.assign(batch_size, VALID_BATCH_SIZE))
115 |     img_batch, label_batch = queueSelection([valid_runner, train_runner], \
116 |                                              tf.cast(is_training, tf.int32), \
117 |                                              batch_size)
118 |     # img_batch, label_batch = doubleQueue(train_runner, valid_runner, \
119 |     #                                      is_training, TRAIN_BATCH_SIZE, \
120 |     #                                      VALID_BATCH_SIZE)
121 | 


--------------------------------------------------------------------------------
/YeNet.py:
--------------------------------------------------------------------------------
  1 | import tensorflow as tf
  2 | from tensorflow.contrib import layers
  3 | from tensorflow.contrib.framework import add_arg_scope, arg_scope, arg_scoped_arguments
  4 | import layers as my_layers
  5 | from utils import *
  6 | 
  7 | SRM_Kernels = np.load('SRM_Kernels.npy')
  8 | 
  9 | class YeNet(Model):
 10 |     def __init__(self, is_training=None, data_format='NCHW', \
 11 |                  with_bn=False, tlu_threshold=3):
 12 |         super(YeNet, self).__init__(is_training=is_training, \
 13 |                                     data_format=data_format)
 14 |         self.with_bn = with_bn
 15 |         self.tlu_threshold = tlu_threshold
 16 | 
 17 |     def _build_model(self, inputs):
 18 |         self.inputs = inputs
 19 |         if self.data_format == 'NCHW':
 20 |             channel_axis = 1
 21 |             _inputs = tf.cast(tf.transpose(inputs, [0, 3, 1, 2]), tf.float32)
 22 |         else:
 23 |             channel_axis = 3
 24 |             _inputs = tf.cast(inputs, tf.float32)
 25 |         self.L = []
 26 |         with arg_scope([layers.avg_pool2d], \
 27 |                 padding='VALID', data_format=self.data_format):
 28 |             with tf.variable_scope('SRM_preprocess'):
 29 |                 W_SRM = tf.get_variable('W', initializer=SRM_Kernels, \
 30 |                             dtype=tf.float32, \
 31 |                             regularizer=None)
 32 |                 b = tf.get_variable('b', shape=[30], dtype=tf.float32, \
 33 |                             initializer=tf.constant_initializer(0.))
 34 |                 self.L.append(tf.nn.bias_add( \
 35 |                         tf.nn.conv2d(_inputs, \
 36 |                         W_SRM, [1,1,1,1], 'VALID', \
 37 |                         data_format=self.data_format), b, \
 38 |                         data_format=self.data_format, name='Layer1'))
 39 |                 self.L.append(tf.clip_by_value(self.L[-1], \
 40 |                               -self.tlu_threshold, self.tlu_threshold, \
 41 |                               name='TLU'))
 42 |             with tf.variable_scope('ConvNetwork'):
 43 |                 with arg_scope([my_layers.conv2d], num_outputs=30, \
 44 |                         kernel_size=3, stride=1, padding='VALID', \
 45 |                         data_format=self.data_format, \
 46 |                         activation_fn=tf.nn.relu, \
 47 |                         weights_initializer=layers.xavier_initializer_conv2d(), \
 48 |                         weights_regularizer=layers.l2_regularizer(5e-4), \
 49 |                         biases_initializer=tf.constant_initializer(0.2), \
 50 |                         biases_regularizer=None), arg_scope([layers.batch_norm], \
 51 |                         decay=0.9, center=True, scale=True, \
 52 |                         updates_collections=None, is_training=self.is_training, \
 53 |                         fused=True, data_format=self.data_format):
 54 |                     if self.with_bn:
 55 |                         self.L.append(layers.batch_norm(self.L[-1], \
 56 |                                       scope='Norm1'))
 57 |                     self.L.append(my_layers.conv2d(self.L[-1], \
 58 |                                   scope='Layer2'))
 59 |                     if self.with_bn:
 60 |                         self.L.append(layers.batch_norm(self.L[-1], \
 61 |                                       scope='Norm2'))
 62 |                     self.L.append(my_layers.conv2d(self.L[-1], \
 63 |                                   scope='Layer3'))
 64 |                     if self.with_bn:
 65 |                         self.L.append(layers.batch_norm(self.L[-1], \
 66 |                                       scope='Norm3'))
 67 |                     self.L.append(my_layers.conv2d(self.L[-1], \
 68 |                                   scope='Layer4'))
 69 |                     if self.with_bn:
 70 |                         self.L.append(layers.batch_norm(self.L[-1], \
 71 |                                       scope='Norm4'))
 72 |                     self.L.append(layers.avg_pool2d(self.L[-1], \
 73 |                                   kernel_size=[2,2], scope='Stride1'))
 74 |                     with arg_scope([my_layers.conv2d], kernel_size=5, \
 75 |                                    num_outputs=32):
 76 |                         self.L.append(my_layers.conv2d(self.L[-1], \
 77 |                                       scope='Layer5'))
 78 |                         if self.with_bn:
 79 |                             self.L.append(layers.batch_norm(self.L[-1], \
 80 |                                           scope='Norm5'))
 81 |                         self.L.append(layers.avg_pool2d(self.L[-1], \
 82 |                                       kernel_size=[3,3], \
 83 |                                       scope='Stride2'))
 84 |                         self.L.append(my_layers.conv2d(self.L[-1], \
 85 |                                       scope='Layer6'))
 86 |                         if self.with_bn:
 87 |                             self.L.append(layers.batch_norm(self.L[-1], \
 88 |                                           scope='Norm6'))
 89 |                         self.L.append(layers.avg_pool2d(self.L[-1], \
 90 |                                       kernel_size=[3,3], \
 91 |                                       scope='Stride3'))
 92 |                         self.L.append(my_layers.conv2d(self.L[-1], \
 93 |                                       scope='Layer7'))
 94 |                         if self.with_bn:
 95 |                             self.L.append(layers.batch_norm(self.L[-1], \
 96 |                                           scope='Norm7'))
 97 |                     self.L.append(layers.avg_pool2d(self.L[-1], \
 98 |                                   kernel_size=[3,3], \
 99 |                                   scope='Stride4'))
100 |                     self.L.append(my_layers.conv2d(self.L[-1], \
101 |                                   num_outputs=16, \
102 |                                   scope='Layer8'))
103 |                     if self.with_bn:
104 |                         self.L.append(layers.batch_norm(self.L[-1], \
105 |                                       scope='Norm8'))
106 |                     self.L.append(my_layers.conv2d(self.L[-1], \
107 |                                   num_outputs=16, stride=3, \
108 |                                   scope='Layer9'))
109 |                     if self.with_bn:
110 |                         self.L.append(layers.batch_norm(self.L[-1], \
111 |                                       scope='Norm9'))
112 |                 self.L.append(layers.flatten(self.L[-1]))
113 |                 self.L.append(layers.fully_connected(self.L[-1], num_outputs=2, \
114 |                         activation_fn=None, normalizer_fn=None, \
115 |                         weights_initializer=tf.random_normal_initializer(mean=0., stddev=0.01), \
116 |                         biases_initializer=tf.constant_initializer(0.), scope='ip'))
117 |         self.outputs = self.L[-1]
118 |         return self.outputs
119 |     
120 | 


--------------------------------------------------------------------------------
/generator.py:
--------------------------------------------------------------------------------
  1 | import numpy as np
  2 | from scipy import misc, io
  3 | from glob import glob
  4 | import random
  5 | from itertools import izip
  6 | from random import random as rand
  7 | from random import shuffle
  8 | import h5py
  9 | 
 10 | def gen_embedding_otf(cover_dir, beta_dir, shuf_pair, \
 11 |                       thread_idx=0, n_threads=1):
 12 |     cover_list = sorted(glob(cover_dir + '/*'))
 13 |     beta_list = sorted(glob(beta_dir + '/*'))
 14 |     nb_data = len(cover_list)
 15 |     assert len(beta_list) != 0, "the beta directory '%s' is empty" % beta_dir
 16 |     assert nb_data != 0, "the cover directory '%s' is empty" % cover_dir
 17 |     assert len(beta_list) == nb_data, "the cover directory and " + \
 18 |                                       "the beta directory don't " + \
 19 |                                       "have the same number of files " + \
 20 |                                       "respectively %d and %d" % (nb_data, \
 21 |                                       len(beta_list))
 22 |     img = misc.imread(cover_list[0])
 23 |     img_shape = img.shape
 24 |     batch = np.empty((2, img_shape[0], img_shape[1], 1), dtype='uint8')
 25 |     beta_map = np.empty(img_shape, dtype='<f8')
 26 |     inf_map = np.empty(img_shape, dtype='bool')
 27 |     rand_arr = np.empty(img_shape, dtype='float64')
 28 |     shuf_cov = np.empty(img_shape, dtype='uint8')
 29 |     while True:
 30 |         if shuf_pair:
 31 |             list_i = np.random.permutation(nb_data)
 32 |             list_j = np.random.permutation(nb_data)
 33 |             for i, j in izip(list_i, list_j):
 34 |                 batch[0,:,:,0] = misc.imread(cover_list[i])
 35 |                 beta_map[:,:] = io.loadmat(beta_list[j])['pChange']
 36 |                 shuf_cov[:,:] = misc.imread(cover_list[j])
 37 |                 rand_arr[:,:] = np.random.rand(img_shape[0], img_shape[1])
 38 |                 inf_map[:,:] = rand_arr < (beta_map / 2.)
 39 |                 batch[1,:,:,0] = np.copy(shuf_cov)
 40 |                 batch[1,np.logical_and(shuf_cov != 255, inf_map),0] += 1
 41 |                 batch[1,np.logical_and(shuf_cov != 0, \
 42 |                       np.logical_and(np.logical_not(inf_map), \
 43 |                       rand_arr < beta_map)), 0] -= 1
 44 |                 rot = random.randint(0,3)
 45 |                 if rand() < 0.5:
 46 |                     yield [np.rot90(batch, rot, axes=[1,2]), np.array([0,1], dtype='uint8')]
 47 |                 else:
 48 |                     yield [np.flip(np.rot90(batch, rot, axes=[1,2]), axis=2), np.array([0,1], dtype='uint8')]
 49 |         else:
 50 |             list_i = np.random.permutation(nb_data)
 51 |             for i in list_i:
 52 |                 batch[0,:,:,0] = misc.imread(cover_list[i])
 53 |                 beta_map[:,:] = io.loadmat(beta_list[i])['pChange']
 54 |                 rand_arr[:,:] = np.random.rand(img_shape[0], img_shape[1])
 55 |                 inf_map[:,:] = rand_arr < (beta_map / 2.)
 56 |                 batch[1,:,:,0] = np.copy(batch[0,:,:,0])
 57 |                 batch[1,np.logical_and(batch[0,:,:,0] != 255, inf_map),0] += 1
 58 |                 batch[1,np.logical_and(batch[0,:,:,0] != 0, \
 59 |                       np.logical_and(np.logical_not(inf_map), \
 60 |                       rand_arr < beta_map)), 0] -= 1
 61 |                 rot = random.randint(0,3)
 62 |                 if rand() < 0.5:
 63 |                     yield [np.rot90(batch, rot, axes=[1,2]), np.array([0,1], dtype='uint8')]
 64 |                 else:
 65 |                     yield [np.flip(np.rot90(batch, rot, axes=[1,2]), axis=2), np.array([0,1], dtype='uint8')]
 66 | 
 67 | def gen_all_flip_and_rot(cover_dir, stego_dir, thread_idx, n_threads):
 68 |     cover_list = sorted(glob(cover_dir + '/*'))
 69 |     stego_list = sorted(glob(stego_dir + '/*'))
 70 |     nb_data = len(cover_list)
 71 |     assert len(stego_list) != 0, "the beta directory '%s' is empty" % stego_dir
 72 |     assert nb_data != 0, "the cover directory '%s' is empty" % cover_dir
 73 |     assert len(stego_list) == nb_data, "the cover directory and " + \
 74 |                                       "the beta directory don't " + \
 75 |                                       "have the same number of files " + \
 76 |                                       "respectively %d and %d" % (nb_data, + \
 77 |                                       len(stego_list))
 78 |     img = misc.imread(cover_list[0])
 79 |     img_shape = img.shape
 80 |     batch = np.empty((2,img_shape[0],img_shape[1],1), dtype='uint8')
 81 |     iterable = zip(cover_list, stego_list)
 82 |     for cover_path, stego_path in iterable:
 83 |         batch[0,:,:,0] = misc.imread(cover_path)
 84 |         batch[1,:,:,0] = misc.imread(stego_path)
 85 |         for rot in range(4):
 86 |             yield [np.rot90(batch, rot, axes=[1,2]), np.array([0,1], dtype='uint8')]
 87 |         for rot in range(4):
 88 |             yield [np.flip(np.rot90(batch, rot, axes=[1,2]), axis=2), np.array([0,1], dtype='uint8')]
 89 | 
 90 | def gen_flip_and_rot(cover_dir, stego_dir, shuf_pair=False, thread_idx=0, n_threads=1):
 91 |     cover_list = sorted(glob(cover_dir + '/*'))
 92 |     stego_list = sorted(glob(stego_dir + '/*'))
 93 |     nb_data = len(cover_list)
 94 |     assert len(stego_list) != 0, "the beta directory '%s' is empty" % stego_dir
 95 |     assert nb_data != 0, "the cover directory '%s' is empty" % cover_dir
 96 |     assert len(stego_list) == nb_data, "the cover directory and " + \
 97 |                                       "the beta directory don't " + \
 98 |                                       "have the same number of files " + \
 99 |                                       "respectively %d and %d" % (nb_data, + \
100 |                                       len(stego_list))
101 |     img = misc.imread(cover_list[0])
102 |     img_shape = img.shape
103 |     batch = np.empty((2,img_shape[0],img_shape[1],1), dtype='uint8')
104 |     if not shuf_pair:
105 |         iterable = zip(cover_list, stego_list)
106 |     while True:
107 |         if shuf_pair:
108 |             shuffle(cover_list)
109 |             shuffle(stego_list)
110 |             iterable = izip(cover_list, stego_list)
111 |         else:
112 |             shuffle(iterable)
113 |         for cover_path, stego_path in iterable:
114 |             batch[0,:,:,0] = misc.imread(cover_path)
115 |             batch[1,:,:,0] = misc.imread(stego_path)
116 |             rot = random.randint(0,3)
117 |             if rand() < 0.5:
118 |                 yield [np.rot90(batch, rot, axes=[1,2]), np.array([0,1], dtype='uint8')]
119 |             else:
120 |                 yield [np.flip(np.rot90(batch, rot, axes=[1,2]), axis=2), np.array([0,1], dtype='uint8')]
121 | 
122 | def gen_valid(cover_dir, stego_dir, thread_idx, n_threads):
123 |     cover_list = sorted(glob(cover_dir + '/*'))
124 |     stego_list = sorted(glob(stego_dir + '/*'))
125 |     nb_data = len(cover_list)
126 |     assert len(stego_list) != 0, "the beta directory '%s' is empty" % stego_dir
127 |     assert nb_data != 0, "the cover directory '%s' is empty" % cover_dir
128 |     assert len(stego_list) == nb_data, "the cover directory and " + \
129 |                                       "the beta directory don't " + \
130 |                                       "have the same number of files " + \
131 |                                       "respectively %d and %d" % (nb_data, \
132 |                                       len(stego_list))
133 |     img = misc.imread(cover_list[0])
134 |     img_shape = img.shape
135 |     batch = np.empty((2,img_shape[0],img_shape[1],1), dtype='uint8')
136 |     labels = np.array([0, 1], dtype='uint8')
137 |     while True:
138 |         for cover_path, stego_path in zip(cover_list, stego_list):
139 |             batch[0,:,:,0] = misc.imread(cover_path)
140 |             batch[1,:,:,0] = misc.imread(stego_path)
141 |             yield [batch, labels]
142 | 
143 | 
144 | # def trainGen(thread_idx, n_threads):
145 | #     batch = np.empty((2,256,256,1), dtype='uint8')
146 | #     beta_map = np.empty((256, 256), dtype='<f8')
147 | #     inf_map = np.empty((256, 256), dtype='bool')
148 | #     rand_arr = np.empty((256, 256), dtype='float64')
149 | 
150 | #     cover_ds_path = LOG_DIR + '/cover_train' + str(thread_idx) +'.txt'
151 | #     shuf_cover_ds_path = LOG_DIR + '/shuf_cover_train' + str(thread_idx) +'.txt'
152 | #     shuf_beta_ds_path = LOG_DIR + '/shuf_beta_train' + str(thread_idx) +'.txt'
153 | #     createSubFile(DS_PATH + 'train_cover.txt', \
154 | #                   cover_ds_path, thread_idx, n_threads)
155 | #     createSubFile(DS_PATH + 'train_beta.txt', \
156 | #                   shuf_beta_ds_path, thread_idx, n_threads)
157 | #     createSubFile(DS_PATH + 'train_cover.txt', \
158 | #                   shuf_cover_ds_path, thread_idx, n_threads)
159 | #     while True:
160 | #         with open(cover_ds_path, 'r') as f_cover, \
161 | #                 open(shuf_cover_ds_path, 'r') as f_shuf_cover, \
162 | #                 open(shuf_beta_ds_path, 'r') as f_shuf_beta:
163 | #             for cov_path, shuf_cov_path, shuf_beta_path in \
164 | #                     izip(f_cover, f_shuf_cover, f_shuf_beta):
165 | #                 batch[0,:,:,0] = misc.imread(str.strip(cov_path))
166 | #                 shuf_cov = misc.imread(str.strip(shuf_cov_path))
167 | #                 beta_map[:,:] = io.loadmat(str.strip(shuf_beta_path))['pChange']
168 | #                 rand_arr[:,:] = np.random.rand(256, 256)
169 | #                 inf_map[:,:] = rand_arr < (beta_map / 2.)
170 | #                 batch[1,:,:,0] = np.copy(shuf_cov)
171 | #                 batch[1,np.logical_and(shuf_cov != 255, inf_map),0] += 1
172 | #                 batch[1,np.logical_and(shuf_cov != 0, \
173 | #                       np.logical_and(np.logical_not(inf_map), rand_arr < beta_map)), 0] -= 1
174 | #                 rot = random.randint(0,3)
175 | #                 if rand() < 0.5:
176 | #                     yield [np.rot90(batch, rot, axes=[1,2]), np.array([0,1], dtype='uint8')]
177 | #                 else:
178 | #                     yield [np.flip(np.rot90(batch, rot, axes=[1,2]), axis=2), np.array([0,1], dtype='uint8')]
179 | 


--------------------------------------------------------------------------------
/utils.py:
--------------------------------------------------------------------------------
  1 | import tensorflow as tf
  2 | import numpy as np
  3 | import sys
  4 | import time
  5 | from glob import glob
  6 | from functools import partial
  7 | import os
  8 | from os.path import expanduser
  9 | home = expanduser("~")
 10 | user = home.split('/')[-1]
 11 | sys.path.append(home + '/tflib/')
 12 | from queues import *
 13 | from generator import *
 14 | 
 15 | def optimistic_restore(session, save_file, \
 16 |                        graph=tf.get_default_graph()):
 17 |     reader = tf.train.NewCheckpointReader(save_file)
 18 |     saved_shapes = reader.get_variable_to_shape_map()
 19 |     var_names = sorted([(var.name, var.name.split(':')[0]) for var in tf.global_variables()
 20 |             if var.name.split(':')[0] in saved_shapes])    
 21 |     restore_vars = []    
 22 |     for var_name, saved_var_name in var_names:            
 23 |         curr_var = graph.get_tensor_by_name(var_name)
 24 |         var_shape = curr_var.get_shape().as_list()
 25 |         if var_shape == saved_shapes[saved_var_name]:
 26 |             restore_vars.append(curr_var)
 27 |     opt_saver = tf.train.Saver(restore_vars)
 28 |     opt_saver.restore(session, save_file)
 29 | 
 30 | class average_summary(object):
 31 |     def __init__(self, variable, name, num_iterations):
 32 |         self.sum_variable = tf.get_variable(name, shape=[], \
 33 |                                 initializer=tf.constant_initializer(0.), \
 34 |                                 dtype='float32', \
 35 |                                 trainable=False, \
 36 |                                 collections=[tf.GraphKeys.LOCAL_VARIABLES])
 37 |         with tf.control_dependencies([variable]):
 38 |             self.increment_op = tf.assign_add(self.sum_variable, variable)
 39 |         self.mean_variable = self.sum_variable / float(num_iterations)
 40 |         self.summary = tf.summary.scalar(name, self.mean_variable)
 41 |         with tf.control_dependencies([self.summary]):
 42 |             self.reset_variable_op = tf.assign(self.sum_variable, 0)
 43 | 
 44 |     def add_summary(self, sess, writer, step):
 45 |         s, _ = sess.run([self.summary, self.reset_variable_op])
 46 |         writer.add_summary(s, step)
 47 |     
 48 | class Model(object):
 49 |     def __init__(self, is_training=None, data_format='NCHW'):
 50 |         self.data_format = data_format
 51 |         if is_training is None:
 52 |             self.is_training = tf.get_variable('is_training', dtype=tf.bool, \
 53 |                                     initializer=tf.constant_initializer(True), \
 54 |                                     trainable=False)
 55 |         else:
 56 |             self.is_training = is_training
 57 | 
 58 |     def _build_model(self, inputs):
 59 |         raise NotImplementedError('Here is your model definition')
 60 | 
 61 |     def _build_losses(self, labels):
 62 |         self.labels = tf.cast(labels, tf.int64)
 63 |         with tf.variable_scope('loss'):
 64 |             oh = tf.one_hot(self.labels, 2)
 65 |             self.loss = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits( \
 66 |                                   labels=oh, logits=self.outputs))
 67 |         with tf.variable_scope('accuracy'):
 68 |             am = tf.argmax(self.outputs, 1)
 69 |             equal = tf.equal(am, self.labels)
 70 |             self.accuracy = tf.reduce_mean(tf.cast(equal, tf.float32))
 71 |         return self.loss, self.accuracy
 72 | 
 73 | def train(model_class, train_gen, valid_gen, train_batch_size, \
 74 |           valid_batch_size, valid_ds_size, optimizer, \
 75 |           train_interval, valid_interval, max_iter, \
 76 |           save_interval, log_path, num_runner_threads=1, \
 77 |           load_path=None):
 78 |     tf.reset_default_graph()
 79 |     train_runner = GeneratorRunner(train_gen, train_batch_size * 10)
 80 |     valid_runner = GeneratorRunner(valid_gen, valid_batch_size * 10)
 81 |     is_training = tf.get_variable('is_training', dtype=tf.bool, \
 82 |                                   initializer=True, trainable=False)
 83 |     if train_batch_size == valid_batch_size:
 84 |         batch_size = train_batch_size
 85 |         disable_training_op = tf.assign(is_training, False)
 86 |         enable_training_op = tf.assign(is_training, True)
 87 |     else:
 88 |         batch_size = tf.get_variable('batch_size', dtype=tf.int32, \
 89 |                                      initializer=train_batch_size, \
 90 |                                      trainable=False, \
 91 |                                      collections=[tf.GraphKeys.LOCAL_VARIABLES])
 92 |         disable_training_op = tf.group(tf.assign(is_training, False), \
 93 |                                 tf.assign(batch_size, valid_batch_size))
 94 |         enable_training_op = tf.group(tf.assign(is_training, True), \
 95 |                                 tf.assign(batch_size, train_batch_size))
 96 |     img_batch, label_batch = queueSelection([valid_runner, train_runner], \
 97 |                                             tf.cast(is_training, tf.int32), \
 98 |                                             batch_size)
 99 |     model = model_class(is_training, 'NCHW')
100 |     model._build_model(img_batch)
101 |     loss, accuracy = model._build_losses(label_batch)
102 |     regularization_losses = tf.get_collection(
103 |                           tf.GraphKeys.REGULARIZATION_LOSSES)
104 |     regularized_loss = tf.add_n([loss] + regularization_losses)
105 |     train_loss_s = average_summary(loss, 'train_loss', train_interval)
106 |     train_accuracy_s = average_summary(accuracy, 'train_accuracy', \
107 |                                        train_interval)
108 |     valid_loss_s = average_summary(loss, 'valid_loss', \
109 |                                    float(valid_ds_size) / float(valid_batch_size))
110 |     valid_accuracy_s = average_summary(accuracy, 'valid_accuracy', \
111 |                                        float(valid_ds_size) / float(valid_batch_size))
112 |     global_step = tf.get_variable('global_step', dtype=tf.int32, shape=[], \
113 |                                   initializer=tf.constant_initializer(0), \
114 |                                   trainable=False)
115 |     minimize_op = optimizer.minimize(regularized_loss, global_step)
116 |     train_op = tf.group(minimize_op, train_loss_s.increment_op, \
117 |                         train_accuracy_s.increment_op)
118 |     increment_valid = tf.group(valid_loss_s.increment_op, \
119 |                                valid_accuracy_s.increment_op)
120 |     init_op = tf.group(tf.global_variables_initializer(), \
121 |                        tf.local_variables_initializer())
122 |     saver = tf.train.Saver(max_to_keep=10000)
123 |     with tf.Session() as sess:
124 |         sess.run(init_op)
125 |         if load_path is not None:
126 |             loader = tf.train.Saver(reshape=True)
127 |             loader.restore(sess, load_path)
128 |         train_runner.start_threads(sess, num_runner_threads)
129 |         valid_runner.start_threads(sess, 1)
130 |         writer = tf.summary.FileWriter(log_path + '/LogFile/', \
131 |                                        sess.graph)
132 |         start = sess.run(global_step)
133 |         sess.run(disable_training_op)
134 |         sess.run([valid_loss_s.reset_variable_op, \
135 |                   valid_accuracy_s.reset_variable_op, \
136 |                   train_loss_s.reset_variable_op, \
137 |                   train_accuracy_s.reset_variable_op])
138 |         _time = time.time()
139 |         for j in range(0, valid_ds_size, valid_batch_size):
140 |             sess.run([increment_valid])
141 |         _acc_val = sess.run(valid_accuracy_s.mean_variable)
142 |         print "validation:", _acc_val, " | ", \
143 |                 "duration:", time.time() - _time, \
144 |                 "seconds long"
145 |         valid_accuracy_s.add_summary(sess, writer, start)
146 |         valid_loss_s.add_summary(sess, writer, start)
147 |         sess.run(enable_training_op)
148 |         print valid_interval
149 |         for i in xrange(start+1, max_iter+1):
150 |             sess.run(train_op)
151 |             if i % train_interval == 0:
152 |                 train_loss_s.add_summary(sess, writer, i)
153 |                 train_accuracy_s.add_summary(sess, writer, i)
154 |             if i % valid_interval == 0:
155 |                 sess.run(disable_training_op)
156 |                 for j in range(0, valid_ds_size, valid_batch_size):
157 |                     sess.run([increment_valid])
158 |                 valid_loss_s.add_summary(sess, writer, i)
159 |                 valid_accuracy_s.add_summary(sess, writer, i)
160 |                 sess.run(enable_training_op)
161 |             if i % save_interval == 0:
162 |                 saver.save(sess, log_path + '/Model_' + str(i) + '.ckpt')
163 | 
164 | def test_dataset(model_class, gen, batch_size, ds_size, load_path):
165 |     tf.reset_default_graph()
166 |     runner = GeneratorRunner(gen, batch_size * 10)
167 |     img_batch, label_batch = runner.get_batched_inputs(batch_size)
168 |     model = model_class(False, 'NCHW')
169 |     model._build_model(img_batch)
170 |     loss, accuracy = model._build_losses(label_batch)
171 |     loss_summary = average_summary(loss, 'loss',  \
172 |                                    float(ds_size) / float(batch_size))
173 |     accuracy_summary = average_summary(accuracy, 'accuracy',  \
174 |                                    float(ds_size) / float(batch_size))
175 |     increment_op = tf.group(loss_summary.increment_op, \
176 |                             accuracy_summary.increment_op)
177 |     global_step = tf.get_variable('global_step', dtype=tf.int32, shape=[], \
178 |                                   initializer=tf.constant_initializer(0), \
179 |                                   trainable=False)
180 |     init_op = tf.group(tf.global_variables_initializer(), \
181 |                        tf.local_variables_initializer())
182 |     saver = tf.train.Saver(max_to_keep=10000)
183 |     with tf.Session() as sess:
184 |         sess.run(init_op)
185 |         saver.restore(sess, load_path)
186 |         runner.start_threads(sess, 1)
187 |         for j in range(0, ds_size, batch_size):
188 |             sess.run(increment_op)
189 |         mean_loss, mean_accuracy = sess.run([loss_summary.mean_variable ,\
190 |                                              accuracy_summary.mean_variable])
191 |     print "Accuracy:", mean_accuracy, " | Loss:", mean_loss
192 | 
193 | def find_best(model_class, valid_gen, test_gen, valid_batch_size, \
194 |               test_batch_size, valid_ds_size, test_ds_size, load_paths):
195 |     tf.reset_default_graph()
196 |     valid_runner = GeneratorRunner(valid_gen, valid_batch_size * 30)
197 |     img_batch, label_batch = valid_runner.get_batched_inputs(valid_batch_size)
198 |     model = model_class(False, 'NCHW')
199 |     model._build_model(img_batch)
200 |     loss, accuracy = model._build_losses(label_batch)
201 |     loss_summary = average_summary(loss, 'loss',  \
202 |                                           float(valid_ds_size) \
203 |                                           / float(valid_batch_size))
204 |     accuracy_summary = average_summary(accuracy, 'accuracy',  \
205 |                                           float(valid_ds_size) \
206 |                                           / float(valid_batch_size))
207 |     increment_op = tf.group(loss_summary.increment_op, \
208 |                             accuracy_summary.increment_op)
209 |     global_step = tf.get_variable('global_step', dtype=tf.int32, shape=[], \
210 |                                   initializer=tf.constant_initializer(0), \
211 |                                   trainable=False)
212 |     init_op = tf.group(tf.global_variables_initializer(), \
213 |                        tf.local_variables_initializer())
214 |     saver = tf.train.Saver(max_to_keep=10000)
215 |     accuracy_arr = []
216 |     loss_arr = []
217 |     print "validation"
218 |     for load_path in load_paths:
219 |         with tf.Session() as sess:
220 |             sess.run(init_op)
221 |             saver.restore(sess, load_path)
222 |             valid_runner.start_threads(sess, 1)
223 |             _time = time.time()
224 |             for j in range(0, valid_ds_size, valid_batch_size):
225 |                 sess.run(increment_op)
226 |             mean_loss, mean_accuracy = sess.run([loss_summary.mean_variable ,\
227 |                                             accuracy_summary.mean_variable])
228 |             accuracy_arr.append(mean_accuracy)
229 |             loss_arr.append(mean_loss)
230 |             print load_path
231 |             print "Accuracy:", accuracy_arr[-1], "| Loss:", loss_arr[-1], \
232 |                     "in", time.time() - _time, "seconds."
233 |     argmax = np.argmax(accuracy_arr)
234 |     print "best savestate:", load_paths[argmax], "with", \
235 |             accuracy_arr[argmax], "accuracy and", loss_arr[argmax], \
236 |             "loss on validation"
237 |     print "test:"
238 |     test_dataset(model_class, test_gen, test_batch_size, test_ds_size, \
239 |                  load_paths[argmax])
240 |     return argmax, accuracy_arr, loss_arr
241 | 
242 | 
243 | def extract_stats_outputs(model_class, gen, batch_size, ds_size, load_path):
244 |     tf.reset_default_graph()
245 |     runner = GeneratorRunner(gen, batch_size * 10)
246 |     img_batch, label_batch = runner.get_batched_inputs(batch_size)
247 |     model = model_class(False, 'NCHW')
248 |     model._build_model(img_batch)
249 |     global_step = tf.get_variable('global_step', dtype=tf.int32, shape=[], \
250 |                                   initializer=tf.constant_initializer(0), \
251 |                                   trainable=False)
252 |     init_op = tf.group(tf.global_variables_initializer(), \
253 |                        tf.local_variables_initializer())
254 |     saver = tf.train.Saver(max_to_keep=10000)
255 |     stats_outputs_arr = np.empty([ds_size, \
256 |         model.stats_outputs.get_shape().as_list()[1]])
257 |     with tf.Session() as sess:
258 |         sess.run(init_op)
259 |         saver.restore(sess, load_path)
260 |         runner.start_threads(sess, 1)
261 |         for j in range(0, ds_size, batch_size):
262 |             stats_outputs_arr[j:j+batch_size] = sess.run(model.stats_outputs)
263 |     return stats_outputs_arr
264 | 
265 | def stats_outputs_all_datasets(model_class, ds_head_dir, payload, \
266 |                                algorithm, load_path, save_dir):
267 |     if not os.path.exists(save_dir):
268 |         os.makedirs(save_dir + '/')
269 |     payload_str = ''.join(str(payload).strip('.'))
270 |     train_ds_size = len(glob(ds_head_dir + '/train/cover/*'))
271 |     valid_ds_size = len(glob(ds_head_dir + '/valid/cover/*'))
272 |     test_ds_size = len(glob(ds_head_dir + '/test/cover/*'))
273 |     train_gen = partial(gen_all_flip_and_rot, ds_head_dir + \
274 |                         '/train/cover/', ds_head_dir + '/train/' + \
275 |                         algorithm + '/payload' + payload_str + '/stego/')
276 |     valid_gen = partial(gen_valid, ds_head_dir + '/valid/cover/', \
277 |                         ds_head_dir + '/valid/' + algorithm + \
278 |                         '/payload' + payload_str + '/stego/')
279 |     test_gen = partial(gen_valid, ds_head_dir + '/test/cover/', \
280 |                       ds_head_dir + '/test/' + algorithm + \
281 |                       '/payload' + payload_str + '/stego/')
282 |     print "train..."
283 |     stats_outputs = extract_stats_outputs(model_class, train_gen, 16, \
284 |                                                 train_ds_size * 2 * 4 * 2, \
285 |                                                 load_path)
286 |     stats_shape = stats_outputs.shape
287 |     stats_outputs = stats_outputs.reshape(train_ds_size, 2, 4, \
288 |                                           2, stats_shape[-1])
289 |     stats_outputs = np.transpose(stats_outputs, axes=[0,3,2,1,4])
290 |     np.save(save_dir + '/train.npy', stats_outputs)
291 |     print "validation..."
292 |     stats_outputs = extract_stats_outputs(model_class, valid_gen, 16, \
293 |                                           valid_ds_size * 2, load_path)
294 |     np.save(save_dir + '/valid.npy', stats_outputs)
295 |     print "test..."
296 |     stats_outputs = extract_stats_outputs(model_class, test_gen, 16, \
297 |                                           test_ds_size * 2, load_path)
298 |     np.save(save_dir + '/test.npy', stats_outputs)
299 | 


--------------------------------------------------------------------------------
/layers.py:
--------------------------------------------------------------------------------
  1 | import tensorflow as tf
  2 | from tensorflow.contrib import layers
  3 | from tensorflow.contrib.framework import add_arg_scope
  4 | 
  5 | @add_arg_scope
  6 | def double_conv2d(ref_half, real_half,
  7 |                   num_outputs,
  8 |                   kernel_size,
  9 |                   stride=1,
 10 |                   padding='SAME',
 11 |                   data_format=None,
 12 |                   rate=1,
 13 |                   activation_fn=tf.nn.relu,
 14 |                   normalizer_fn=None,
 15 |                   normalize_after_activation=True,
 16 |                   normalizer_params=None,
 17 |                   weights_initializer=layers.xavier_initializer(),
 18 |                   weights_regularizer=None,
 19 |                   biases_initializer=tf.zeros_initializer(),
 20 |                   biases_regularizer=None,
 21 |                   reuse=None,
 22 |                   variables_collections=None,
 23 |                   outputs_collections=None,
 24 |                   trainable=True,
 25 |                   scope=None):
 26 |     with tf.variable_scope(scope, 'Conv', reuse=reuse):
 27 |         if data_format == 'NHWC':
 28 |             num_inputs = real_half.get_shape().as_list()[3]
 29 |             height = real_half.get_shape().as_list()[1]
 30 |             width = real_half.get_shape().as_list()[2]
 31 |             if isinstance(stride, int):
 32 |                 strides = [1, stride, stride, 1]
 33 |             elif isinstance(stride, list) or isinstance(stride, tuple):
 34 |                 if len(stride) == 1:
 35 |                     strides = [1] + stride * 2 + [1]
 36 |                 else:
 37 |                     strides = [1, stride[0], stride[1], 1]
 38 |             else:
 39 |                 raise TypeError('stride is not an int, list or' \
 40 |                                 + 'a tuple, is %s' % type(stride))
 41 |         else:
 42 |             num_inputs = real_half.get_shape().as_list()[1]
 43 |             height = real_half.get_shape().as_list()[2]
 44 |             width = real_half.get_shape().as_list()[3]
 45 |             if isinstance(stride, int):
 46 |                 strides = [1, 1, stride, stride]
 47 |             elif isinstance(stride, list) or isinstance(stride, tuple):
 48 |                 if len(stride) == 1:
 49 |                     strides = [1, 1] + stride * 2
 50 |                 else:
 51 |                     strides = [1, 1, stride[0], stride[1]]
 52 |             else:
 53 |                 raise TypeError('stride is not an int, list or' \
 54 |                                 + 'a tuple, is %s' % type(stride))
 55 |         if isinstance(kernel_size, int):
 56 |             kernel_height = kernel_size
 57 |             kernel_width = kernel_size
 58 |         elif isinstance(kernel_size, list) \
 59 |                 or isinstance(kernel_size, tuple):
 60 |             kernel_height = kernel_size[0]
 61 |             kernel_width = kernel_size[1]
 62 |         else:
 63 |             raise ValueError('kernel_size is not an int, list or' \
 64 |                              + 'a tuple, is %s' % type(kernel_size))
 65 |         weights = tf.get_variable('weights', [kernel_height, \
 66 |                                   kernel_width, num_inputs, num_outputs], \
 67 |                                   'float32', weights_initializer, \
 68 |                                   weights_regularizer, trainable, \
 69 |                                   variables_collections)
 70 |         ref_outputs = tf.nn.conv2d(ref_half, weights, strides, padding, \
 71 |                                   data_format=data_format)
 72 |         real_outputs = tf.nn.conv2d(real_half, weights, strides, padding, \
 73 |                                     data_format=data_format)
 74 |         if biases_initializer is not None:
 75 |             biases = tf.get_variable('biases', [num_outputs], 'float32', \
 76 |                                      biases_initializer, \
 77 |                                      biases_regularizer, \
 78 |                                      trainable, variables_collections)
 79 |             ref_outputs = tf.nn.bias_add(ref_outputs, biases, data_format)
 80 |             real_outputs = tf.nn.bias_add(real_outputs, biases, data_format)
 81 |         if normalizer_fn is not None \
 82 |                 and not normalize_after_activation:
 83 |             normalizer_params = normalizer_params or {}
 84 |             ref_outputs, real_outputs = normalizer_fn(ref_outputs, \
 85 |                                                       real_outputs,  \
 86 |                                                       **normalizer_params)
 87 |         if activation_fn is not None:
 88 |             ref_outputs = activation_fn(ref_outputs)
 89 |             real_outputs = activation_fn(real_outputs)
 90 |         if normalizer_fn is not None and normalize_after_activation:
 91 |             normalizer_params = normalizer_params or {}
 92 |             ref_outputs, real_outputs = normalizer_fn(ref_outputs, \
 93 |                                                       real_outputs,\
 94 |                                                       **normalizer_params)
 95 |     return ref_outputs, real_outputs
 96 | 
 97 | @add_arg_scope
 98 | def conv2d(inputs,
 99 |            num_outputs,
100 |            kernel_size,
101 |            stride=1,
102 |            padding='SAME',
103 |            data_format=None,
104 |            rate=1,
105 |            activation_fn=tf.nn.relu,
106 |            normalizer_fn=None,
107 |            normalize_after_activation=True,
108 |            normalizer_params=None,
109 |            weights_initializer=layers.xavier_initializer(),
110 |            weights_regularizer=None,
111 |            biases_initializer=tf.zeros_initializer(),
112 |            biases_regularizer=None,
113 |            reuse=None,
114 |            variables_collections=None,
115 |            outputs_collections=None,
116 |            trainable=True,
117 |            scope=None):
118 |     with tf.variable_scope(scope, 'Conv', reuse=reuse):
119 |         if data_format == 'NHWC':
120 |             num_inputs = inputs.get_shape().as_list()[3]
121 |             height = inputs.get_shape().as_list()[1]
122 |             width = inputs.get_shape().as_list()[2]
123 |             if isinstance(stride, int):
124 |                 strides = [1, stride, stride, 1]
125 |             elif isinstance(stride, list) or isinstance(stride, tuple):
126 |                 if len(stride) == 1:
127 |                     strides = [1] + stride * 2 + [1]
128 |                 else:
129 |                     strides = [1, stride[0], stride[1], 1]
130 |             else:
131 |                 raise TypeError('stride is not an int, list or' \
132 |                                 + 'a tuple, is %s' % type(stride))
133 |         else:
134 |             num_inputs = inputs.get_shape().as_list()[1]
135 |             height = inputs.get_shape().as_list()[2]
136 |             width = inputs.get_shape().as_list()[3]
137 |             if isinstance(stride, int):
138 |                 strides = [1, 1, stride, stride]
139 |             elif isinstance(stride, list) or isinstance(stride, tuple):
140 |                 if len(stride) == 1:
141 |                     strides = [1, 1] + stride * 2
142 |                 else:
143 |                     strides = [1, 1, stride[0], stride[1]]
144 |             else:
145 |                 raise TypeError('stride is not an int, list or' \
146 |                                 + 'a tuple, is %s' % type(stride))
147 |         if isinstance(kernel_size, int):
148 |             kernel_height = kernel_size
149 |             kernel_width = kernel_size
150 |         elif isinstance(kernel_size, list) \
151 |                 or isinstance(kernel_size, tuple):
152 |             kernel_height = kernel_size[0]
153 |             kernel_width = kernel_size[1]
154 |         else:
155 |             raise ValueError('kernel_size is not an int, list or' \
156 |                              + 'a tuple, is %s' % type(kernel_size))
157 |         weights = tf.get_variable('weights', [kernel_height, \
158 |                                   kernel_width, num_inputs, num_outputs], \
159 |                                   'float32', weights_initializer, \
160 |                                   weights_regularizer, trainable, \
161 |                                   variables_collections)
162 |         outputs = tf.nn.conv2d(inputs, weights, strides, padding, \
163 |                             data_format=data_format)
164 |         if biases_initializer is not None:
165 |             biases = tf.get_variable('biases', [num_outputs], 'float32', \
166 |                                      biases_initializer, \
167 |                                      biases_regularizer, \
168 |                                      trainable, variables_collections)
169 |             outputs = tf.nn.bias_add(outputs, biases, data_format)
170 |         if normalizer_fn is not None \
171 |                 and not normalize_after_activation:
172 |             normalizer_params = normalizer_params or {}
173 |             outputs = normalizer_fn(outputs, **normalizer_params)
174 |         if activation_fn is not None:
175 |             outputs = activation_fn(outputs)
176 |         if normalizer_fn is not None and normalize_after_activation:
177 |             normalizer_params = normalizer_params or {}
178 |             outputs = normalizer_fn(outputs, **normalizer_params)
179 |     return outputs
180 | 
181 | class Vbn_double(object):
182 |     def __init__(self, x, epsilon=1e-5, scope=None):
183 |         shape = x.get_shape().as_list()
184 |         needs_reshape = len(shape) != 4
185 |         if needs_reshape:
186 |             orig_shape = shape
187 |             if len(shape) == 2:
188 |                 if data_format == 'NCHW':
189 |                     x = tf.reshape(x, [shape[0], shape[1], 0, 0])
190 |                 else:
191 |                     x = tf.reshape(x, [shape[0], 1, 1, shape[1]])
192 |             elif len(shape) == 1:
193 |                 x = tf.reshape(x, [shape[0], 1, 1, 1])
194 |             else:
195 |                 assert False, shape
196 |             shape = x.get_shape().as_list()
197 |         with tf.variable_scope(scope):
198 |             self.epsilon = epsilon
199 |             self.scope = scope
200 |             self.mean, self.var = tf.nn.moments(x, [0,2,3], \
201 |                                                 keep_dims=True)
202 |             self.inv_std = tf.rsqrt(self.var + epsilon)
203 |             self.batch_size = int(x.get_shape()[0])
204 |             out = self._normalize(x, self.mean, self.inv_std)
205 |             if needs_reshape:
206 |                 out = tf.reshape(out, orig_shape)
207 |             self.reference_output = out
208 | 
209 |     def __call__(self, x):
210 |         shape = x.get_shape().as_list()
211 |         needs_reshape = len(shape) != 4
212 |         if needs_reshape:
213 |             orig_shape = shape
214 |             if len(shape) == 2:
215 |                 if self.data_format == 'NCHW':
216 |                     x = tf.reshape(x, [shape[0], shape[1], 0, 0])
217 |                 else:
218 |                     x = tf.reshape(x, [shape[0], 1, 1, shape[1]])
219 |             elif len(shape) == 1:
220 |                 x = tf.reshape(x, [shape[0], 1, 1, 1])
221 |             else:
222 |                 assert False, shape
223 |         with tf.variable_scope(self.scope, reuse=True):
224 |             out = self._normalize(x, self.mean, self.inv_std)
225 |             if needs_reshape:
226 |                 out = tf.reshape(out, orig_shape)
227 |         return out
228 | 
229 |     def _normalize(self, x, mean, inv_std):
230 |         shape = x.get_shape().as_list()
231 |         assert len(shape) == 4
232 |         gamma = tf.get_variable("gamma", [1,shape[1],1,1],
233 |                         initializer=tf.constant_initializer(1.))
234 |         beta = tf.get_variable("beta", [1,shape[1],1,1],
235 |                         initializer=tf.constant_initializer(0.))
236 |         coeff = gamma * inv_std
237 |         return (x * coeff) + (beta - mean * coeff)
238 | 
239 | @add_arg_scope
240 | def vbn_double(ref_half, real_half, center=True, scale=True, epsilon=1e-5, \
241 |                data_format='NCHW', instance_norm=True, scope=None, \
242 |                reuse=None):
243 |     assert isinstance(epsilon, float)
244 |     shape = real_half.get_shape().as_list()
245 |     batch_size = int(real_half.get_shape()[0])
246 |     with tf.variable_scope(scope, 'VBN', reuse=reuse):
247 |         if data_format == 'NCHW':
248 |             if scale:
249 |                 gamma = tf.get_variable("gamma", [1,shape[1],1,1],
250 |                             initializer=tf.constant_initializer(1.))
251 |             if center:
252 |                 beta = tf.get_variable("beta", [1,shape[1],1,1],
253 |                             initializer=tf.constant_initializer(0.))
254 |             ref_mean, ref_var = tf.nn.moments(ref_half, [0,2,3], \
255 |                                               keep_dims=True)
256 |         else:
257 |             if scale:
258 |                 gamma = tf.get_variable("gamma", [1,1,1,shape[-1]],
259 |                             initializer=tf.constant_initializer(1.))
260 |             if center:
261 |                 beta = tf.get_variable("beta", [1,1,1,shape[-1]],
262 |                             initializer=tf.constant_initializer(0.))
263 |             ref_mean, ref_var = tf.nn.moments(ref_half, [0,1,2], \
264 |                                               keep_dims=True)
265 |         def _normalize(x, mean, var):
266 |             inv_std = tf.rsqrt(var + epsilon)
267 |             if scale:
268 |                 coeff = inv_std * gamma
269 |             else:
270 |                 coeff = inv_std
271 |             if center:
272 |                 return (x * coeff) + (beta - mean * coeff)
273 |             else:
274 |                 return (x - mean) * coeff
275 |         if instance_norm:
276 |             if data_format == 'NCHW':
277 |                 real_mean, real_var = tf.nn.moments(real_half, [2,3], \
278 |                                                   keep_dims=True)
279 |             else:
280 |                 real_mean, real_var = tf.nn.moments(real_half, [1,2], \
281 |                                                   keep_dims=True)
282 |             real_coeff = 1. / (batch_size + 1.)
283 |             ref_coeff = 1. - real_coeff
284 |             new_mean = real_coeff * real_mean + ref_coeff * ref_mean
285 |             new_var = real_coeff * real_var + ref_coeff * ref_var
286 |             ref_output = _normalize(ref_half, ref_mean, ref_var)
287 |             real_output = _normalize(real_half, new_mean, new_var)
288 |         else:
289 |             ref_output = _normalize(ref_half, ref_mean, ref_var)
290 |             real_output = _normalize(real_half, ref_mean, ref_var)
291 |         return ref_output, real_output
292 | 
293 | 
294 | @add_arg_scope
295 | def vbn_single(x, center=True, scale=True, \
296 |                epsilon=1e-5, data_format='NCHW', \
297 |                instance_norm=True, scope=None, \
298 |                reuse=None):
299 |     assert isinstance(epsilon, float)
300 |     shape = x.get_shape().as_list()
301 |     if shape[0] is None:
302 |         half_size = x.shape[0] // 2
303 |     else:
304 |         half_size = shape[0] // 2
305 |     needs_reshape = len(shape) != 4
306 |     if needs_reshape:
307 |         orig_shape = shape
308 |         if len(shape) == 2:
309 |             if data_format == 'NCHW':
310 |                 x = tf.reshape(x, [shape[0], shape[1], 0, 0])
311 |             else:
312 |                 x = tf.reshape(x, [shape[0], 1, 1, shape[1]])
313 |         elif len(shape) == 1:
314 |             x = tf.reshape(x, [shape[0], 1, 1, 1])
315 |         else:
316 |             assert False, shape
317 |         shape = x.get_shape().as_list()
318 |     batch_size = int(x.get_shape()[0])
319 |     with tf.variable_scope(scope, 'VBN', reuse=reuse):
320 |         ref_half = tf.slice(x, [0,0,0,0], [half_size, shape[1], \
321 |                             shape[2], shape[3]])
322 |         if data_format == 'NCHW':
323 |             if scale:
324 |                 gamma = tf.get_variable("gamma", [1,shape[1],1,1],
325 |                             initializer=tf.constant_initializer(1.))
326 |             if center:
327 |                 beta = tf.get_variable("beta", [1,shape[1],1,1],
328 |                             initializer=tf.constant_initializer(0.))
329 |             ref_mean, ref_var = tf.nn.moments(ref_half, [0,2,3], \
330 |                                               keep_dims=True)
331 |         else:
332 |             if scale:
333 |                 gamma = tf.get_variable("gamma", [1,1,1,shape[-1]],
334 |                             initializer=tf.constant_initializer(1.))
335 |             if center:
336 |                 beta = tf.get_variable("beta", [1,1,1,shape[-1]],
337 |                             initializer=tf.constant_initializer(0.))
338 |             ref_mean, ref_var = tf.nn.moments(ref_half, [0,1,2], \
339 |                                               keep_dims=True)
340 |         def _normalize(x, mean, var):
341 |             inv_std = tf.rsqrt(var + epsilon)
342 |             if scale:
343 |                 coeff = inv_std * gamma
344 |             else:
345 |                 coeff = inv_std
346 |             if center:
347 |                 return (x * coeff) + (beta - mean * coeff)
348 |             else:
349 |                 return (x - mean) * coeff
350 |         if instance_norm:
351 |             real_half = tf.slice(x, [half_size,0,0,0], \
352 |                                  [half_size, shape[1], shape[2], shape[3]])
353 |             if data_format == 'NCHW':
354 |                 real_mean, real_var = tf.nn.moments(real_half, [2,3], \
355 |                                                   keep_dims=True)
356 |             else:
357 |                 real_mean, real_var = tf.nn.moments(real_half, [1,2], \
358 |                                                   keep_dims=True)
359 |             real_coeff = 1. / (batch_size + 1.)
360 |             ref_coeff = 1. - real_coeff
361 |             new_mean = real_coeff * real_mean + ref_coeff * ref_mean
362 |             new_var = real_coeff * real_var + ref_coeff * ref_var
363 |             ref_output = _normalize(ref_half, ref_mean, ref_var)
364 |             real_output = _normalize(real_half, new_mean, new_var)
365 |             return tf.concat([ref_output, real_output], axis=0)
366 |         else:
367 |             return _normalize(x, ref_mean, ref_var)
368 | 
369 | 


--------------------------------------------------------------------------------