def step_up(name,
bottom_input,
side_input,
filter_size=3,
res_blocks=2,
keep_prob=1.,
training=False):
with tf.variable_scope(name):
added_input = layers.upconv_add_block(bottom_input,
side_input,
data_format="NCHW")
conv_out = added_input
for i in xrange(res_blocks):
with tf.variable_scope("res_block_" + str(i)):
conv_out = layers.res_block(conv_out,
filter_size,
channel_multiplier=1,
depthwise_multiplier=2,
convolutions=2,
training=training,
activation=activation,
batch_norm=batch_norm,
data_format="NCHW")
result = layers.dropout(conv_out, keep_prob)
return result
def step_down(name, input_, filter_size=3, res_blocks=2, keep_prob=1., training=False):
with tf.variable_scope(name):
with tf.variable_scope("res_block_0"):
conv_out, tiled_input = layers.res_block(input_, filter_size, channel_multiplier=2, depthwise_multiplier=2, convolutions=2, training=training, activation=activation, batch_norm=batch_norm, data_format="NCHW")
for i in xrange(1, res_blocks):
with tf.variable_scope("res_block_" + str(i)):
conv_out = layers.res_block(conv_out, filter_size, channel_multiplier=1, depthwise_multiplier=2, convolutions=2, training=training, activation=activation, batch_norm=batch_norm, data_format="NCHW")
conv_out = conv_out + tiled_input
pool_out = layers.max_pool(conv_out, pool_size, data_format="NCHW")
bottom_out = layers.dropout(pool_out, keep_prob)
side_out = layers.dropout(conv_out, keep_prob)
return bottom_out, side_out
def generator(real_img, desired_au, reuse=False):
'''
:param:
real_img: RGB face images, shape [batch, 128, 128, 3], value [-1,1].
desired_au: AU value, shape [batch, 17], value [0,1].
:return:
fake_img: RGB generate face, shape [batch, 128, 128, 3], value [-1,1].
fake_mask: face mask, shape [batch, 128, 128, 1], value [0,1].
'''
with tf.variable_scope('generator') as scope:
if reuse:
scope.reuse_variables()
desired_au = tf.expand_dims(desired_au, axis=1, name='ExpandDims1')
desired_au = tf.expand_dims(desired_au, axis=2, name='ExpandDims2')
desired_au = tf.tile(desired_au, multiples=[1,128,128,1], name='Tile')
x = tf.concat([real_img, desired_au], axis=3, name='Concat')
x = conv2d(x, out_channels=64, kernel_size=7, strides=1, name='Conv1')
x = instance_norm(x, name='InstNorm1')
x = tf.nn.relu(x, name='ReLU1')
x = conv2d(x, out_channels=128, kernel_size=4, strides=2, name='Conv2')
x = instance_norm(x, name='InstNorm2')
x = tf.nn.relu(x, name='ReLU2')
x = conv2d(x, out_channels=256, kernel_size=4, strides=2, name='Conv3')
x = instance_norm(x, name='InstNorm3')
x = tf.nn.relu(x, name='ReLU3')
for i in range(1, 7):
x = res_block(x, out_channels=256, name='ResBlock'+str(i))
x = deconv2d(x, out_channels=128, kernel_size=4, stride=2, name='Deconv1')
x = instance_norm(x, name='InstNorm4')
x = tf.nn.relu(x, name='ReLU4')
x = deconv2d(x, out_channels=64, kernel_size=4, stride=2, name='Deconv2')
x = instance_norm(x, name='InstNorm5')
features = tf.nn.relu(x, name='ReLU5')
x = conv2d(features, out_channels=3, kernel_size=7, strides=1, name='ConvImg')
fake_img = tf.tanh(x, name='Tanh')
x = conv2d(features, out_channels=1, kernel_size=7, strides=1, name='ConvMask')
fake_mask = tf.sigmoid(x, name='Sigmoid')
return fake_img, fake_mask
def _network(self, inputs):
init = tf.initializers.variance_scaling(scale=2, mode='fan_in')
reg = tf.contrib.layers.l2_regularizer(1.0)
with tf.variable_scope('preprocessing'):
layers = [(0, ly.casual_conv1d(inputs,
kernel_size=self.kernel_size,
filters=self.res_filters,
dilation=1,
activation=None))]
with tf.variable_scope('residual_stack'):
max_dilation_exp = int(np.log2(self.max_dilation)) + 1
for i in range(self.num_layers):
dilation = 2 ** (i % max_dilation_exp)
layers.append(ly.res_block(
layers[-1][1],
kernel_size=self.kernel_size,
dilation=dilation,
skip_filters=self.skip_filters,
dilation_filters=self.dilation_filters,
res_filters=self.res_filters))
with tf.variable_scope('postprocessing'):
skips = [l[0] for l in layers[1:]]
skips = tf.add_n(skips)
skips = tf.nn.relu(skips)
skips = tf.layers.conv1d(skips,
filters=self.skip_filters,
kernel_size=1,
padding='same',
activation=tf.nn.relu,
kernel_initializer=init,
kernel_regularizer=reg)
skips = tf.layers.conv1d(skips,
filters=256,
kernel_size=1,
padding='same',
kernel_initializer=init,
kernel_regularizer=reg)
return skips
def parameter_efficient(in_channels=1, out_channels=2, start_filters=64, input_side_length=256, depth=4, res_blocks=2, filter_size=3, sparse_labels=True, batch_size=1, activation="cReLU", batch_norm=True):
activation = str.lower(activation)
if activation not in ["relu", "crelu"]:
raise ValueError("activation must be \"ReLU\" or \"cReLU\".")
pool_size = 2
# Define inputs and helper functions #
with tf.variable_scope('inputs'):
inputs = tf.placeholder(tf.float32, shape=(batch_size, input_side_length, input_side_length, in_channels), name='inputs')
if sparse_labels:
ground_truth = tf.placeholder(tf.int32, shape=(batch_size, input_side_length, input_side_length), name='labels')
else:
ground_truth = tf.placeholder(tf.float32, shape=(batch_size, input_side_length, input_side_length, out_channels), name='labels')
keep_prob = tf.placeholder(tf.float32, shape=[], name='keep_prob')
training = tf.placeholder(tf.bool, shape=[], name="training")
network_input = tf.transpose(inputs, perm=[0, 3, 1, 2])
# [conv -> conv -> max pool -> drop out] + parameter updates
def step_down(name, input_, filter_size=3, res_blocks=2, keep_prob=1., training=False):
with tf.variable_scope(name):
with tf.variable_scope("res_block_0"):
conv_out, tiled_input = layers.res_block(input_, filter_size, channel_multiplier=2, depthwise_multiplier=2, convolutions=2, training=training, activation=activation, batch_norm=batch_norm, data_format="NCHW")
for i in xrange(1, res_blocks):
with tf.variable_scope("res_block_" + str(i)):
conv_out = layers.res_block(conv_out, filter_size, channel_multiplier=1, depthwise_multiplier=2, convolutions=2, training=training, activation=activation, batch_norm=batch_norm, data_format="NCHW")
conv_out = conv_out + tiled_input
pool_out = layers.max_pool(conv_out, pool_size, data_format="NCHW")
bottom_out = layers.dropout(pool_out, keep_prob)
side_out = layers.dropout(conv_out, keep_prob)
return bottom_out, side_out
# parameter updates + [upconv and concat -> drop out -> conv -> conv]
def step_up(name, bottom_input, side_input, filter_size=3, res_blocks=2, keep_prob=1., training=False):
with tf.variable_scope(name):
added_input = layers.upconv_add_block(bottom_input, side_input, data_format="NCHW")
conv_out = added_input
for i in xrange(res_blocks):
with tf.variable_scope("res_block_" + str(i)):
conv_out = layers.res_block(conv_out, filter_size, channel_multiplier=1, depthwise_multiplier=2, convolutions=2, training=training, activation=activation, batch_norm=batch_norm, data_format="NCHW")
result = layers.dropout(conv_out, keep_prob)
return result
# Build the network #
with tf.variable_scope('contracting'):
outputs = []
with tf.variable_scope("step_0"):
# Conv 1
in_filters = in_channels
out_filters = start_filters
stddev = np.sqrt(2. / (filter_size**2 * in_filters))
w = layers.weight_variable([filter_size, filter_size, in_filters, out_filters], stddev=stddev, name="weights")
out_ = tf.nn.conv2d(network_input, w, [1, 1, 1, 1], padding="SAME", data_format="NCHW")
out_ = out_ + layers.bias_variable([out_filters, 1, 1], name='biases')
# Batch Norm 1
if batch_norm:
out_ = tf.layers.batch_normalization(out_, axis=1, momentum=0.999, center=True, scale=True, training=training, trainable=True, name="batch_norm", fused=True)
in_filters = out_filters
# concatenated ReLU
if activation == "crelu":
out_ = tf.concat([out_, -out_], axis=1)
in_filters = 2 * in_filters
out_ = tf.nn.relu(out_)
# Conv 2
stddev = np.sqrt(2. / (filter_size**2 * in_filters))
w = layers.weight_variable([filter_size, filter_size, in_filters, out_filters], stddev=stddev, name="weights")
out_ = tf.nn.conv2d(out_, w, [1, 1, 1, 1], padding="SAME", data_format="NCHW")
out_ = out_ + layers.bias_variable([out_filters, 1, 1], name='biases')
# Res Block 1
conv_out = layers.res_block(out_, filter_size, channel_multiplier=1, depthwise_multiplier=2, convolutions=2, training=training, activation=activation, batch_norm=batch_norm, data_format="NCHW")
pool_out = layers.max_pool(conv_out, pool_size, data_format="NCHW")
bottom_out = layers.dropout(pool_out, keep_prob)
side_out = layers.dropout(conv_out, keep_prob)
outputs.append(side_out)
# Build contracting path
for i in xrange(1, depth):
bottom_out, side_out = step_down('step_' + str(i), bottom_out, filter_size=filter_size, res_blocks=res_blocks, keep_prob=keep_prob, training=training)
outputs.append(side_out)
# Bottom [conv -> conv]
with tf.variable_scope('step_' + str(depth)):
with tf.variable_scope("res_block_0"):
conv_out, tiled_input = layers.res_block(bottom_out, filter_size, channel_multiplier=2, depthwise_multiplier=2, convolutions=2, training=training, activation=activation, batch_norm=batch_norm, data_format="NCHW")
for i in xrange(1, res_blocks):
with tf.variable_scope("res_block_" + str(i)):
conv_out = layers.res_block(conv_out, filter_size, channel_multiplier=1, depthwise_multiplier=2, convolutions=2, training=training, activation=activation, batch_norm=batch_norm, data_format="NCHW")
conv_out = conv_out + tiled_input
current_tensor = layers.dropout(conv_out, keep_prob)
with tf.variable_scope('expanding'):
# Set initial parameter
outputs.reverse()
# Build expanding path
for i in xrange(depth):
current_tensor = step_up('step_' + str(depth + i + 1), current_tensor, outputs[i], filter_size=filter_size, res_blocks=res_blocks, keep_prob=keep_prob, training=training)
# Last layer is a 1x1 convolution to get the predictions
# We don't want an activation function for this one (softmax will be applied later), so we're doing it manually
in_filters = current_tensor.shape.as_list()[1]
stddev = np.sqrt(2. / in_filters)
with tf.variable_scope('classification'):
w = layers.weight_variable([1, 1, in_filters, out_channels], stddev, name='weights')
b = layers.bias_variable([out_channels, 1, 1], name='biases')
conv = tf.nn.conv2d(current_tensor, w, strides=[1, 1, 1, 1], padding="SAME", data_format="NCHW", name='conv')
logits = conv + b
logits = tf.transpose(logits, perm=[0, 2, 3, 1])
return inputs, logits, ground_truth, keep_prob, training
def parameter_efficient(in_channels=1,
out_channels=2,
start_filters=64,
input_side_length=256,
depth=4,
res_blocks=2,
filter_size=3,
sparse_labels=True,
batch_size=1,
activation="cReLU",
batch_norm=True):
"""
Creates the graph for the parameter efficient variant of the U-Net and sets up the appropriate input and output placeholder.
Parameters
----------
in_channels: int
The depth of the input.
out_channels: int
The depth of number of classes of the output.
start_filters : int
The number of filters in the first convolution.
input_side_length: int
The side length of the square input.
depth: int
The depth of the U-part of the network. This is equal to the number of max-pooling layers.
res_blocks: int
The number of residual blocks in between max-pooling layers on the down-path and in between up-convolutions on the up-path.
filter_size: int
The width and height of the filter. The receptive field.
sparse_labels: bool
If true, the labels are integers, one integer per pixel, denoting the class that that pixel belongs to. If false, labels are one-hot encoded.
batch_size: int
The training batch size.
activation: string
Either "ReLU" for the standard ReLU activation or "cReLU" for the concatenated ReLU activation function.
batch_norm: bool
Whether to use batch normalization or not.
Returns
-------
inputs : TF tensor
The network input.
logits: TF tensor
The network output before SoftMax.
ground_truth: TF tensor
The desired output from the ground truth.
keep_prob: TF float
The TF variable holding the keep probability for drop out layers.
training_bool: TF bool
The TF variable holding the boolean value, which switches batch normalization to training or inference mode.
"""
activation = str.lower(activation)
if activation not in ["relu", "crelu"]:
raise ValueError("activation must be \"ReLU\" or \"cReLU\".")
pool_size = 2
# Define inputs and helper functions #
with tf.variable_scope('inputs'):
inputs = tf.placeholder(tf.float32,
shape=(batch_size, input_side_length,
input_side_length, in_channels),
name='inputs')
if sparse_labels:
ground_truth = tf.placeholder(tf.int32,
shape=(batch_size, input_side_length,
input_side_length),
name='labels')
else:
ground_truth = tf.placeholder(tf.float32,
shape=(batch_size, input_side_length,
input_side_length,
out_channels),
name='labels')
keep_prob = tf.placeholder(tf.float32, shape=[], name='keep_prob')
training = tf.placeholder(tf.bool, shape=[], name="training")
network_input = tf.transpose(inputs, perm=[0, 3, 1, 2])
# [conv -> conv -> max pool -> drop out] + parameter updates
def step_down(name,
input_,
filter_size=3,
res_blocks=2,
keep_prob=1.,
training=False):
with tf.variable_scope(name):
with tf.variable_scope("res_block_0"):
conv_out, tiled_input = layers.res_block(
input_,
filter_size,
channel_multiplier=2,
depthwise_multiplier=2,
convolutions=2,
training=training,
activation=activation,
batch_norm=batch_norm,
data_format="NCHW")
for i in xrange(1, res_blocks):
with tf.variable_scope("res_block_" + str(i)):
conv_out = layers.res_block(conv_out,
filter_size,
channel_multiplier=1,
depthwise_multiplier=2,
convolutions=2,
training=training,
activation=activation,
batch_norm=batch_norm,
data_format="NCHW")
conv_out = conv_out + tiled_input
pool_out = layers.max_pool(conv_out, pool_size, data_format="NCHW")
bottom_out = layers.dropout(pool_out, keep_prob)
side_out = layers.dropout(conv_out, keep_prob)
return bottom_out, side_out
# parameter updates + [upconv and concat -> drop out -> conv -> conv]
def step_up(name,
bottom_input,
side_input,
filter_size=3,
res_blocks=2,
keep_prob=1.,
training=False):
with tf.variable_scope(name):
added_input = layers.upconv_add_block(bottom_input,
side_input,
data_format="NCHW")
conv_out = added_input
for i in xrange(res_blocks):
with tf.variable_scope("res_block_" + str(i)):
conv_out = layers.res_block(conv_out,
filter_size,
channel_multiplier=1,
depthwise_multiplier=2,
convolutions=2,
training=training,
activation=activation,
batch_norm=batch_norm,
data_format="NCHW")
result = layers.dropout(conv_out, keep_prob)
return result
# Build the network #
with tf.variable_scope('contracting'):
outputs = []
with tf.variable_scope("step_0"):
# Conv 1
in_filters = in_channels
out_filters = start_filters
stddev = np.sqrt(2. / (filter_size**2 * in_filters))
w = layers.weight_variable(
[filter_size, filter_size, in_filters, out_filters],
stddev=stddev,
name="weights")
out_ = tf.nn.conv2d(network_input,
w, [1, 1, 1, 1],
padding="SAME",
data_format="NCHW")
out_ = out_ + layers.bias_variable([out_filters, 1, 1],
name='biases')
# Batch Norm 1
if batch_norm:
out_ = tf.layers.batch_normalization(out_,
axis=1,
momentum=0.999,
center=True,
scale=True,
training=training,
trainable=True,
name="batch_norm",
fused=True)
in_filters = out_filters
# concatenated ReLU
if activation == "crelu":
out_ = tf.concat([out_, -out_], axis=1)
in_filters = 2 * in_filters
out_ = tf.nn.relu(out_)
# Conv 2
stddev = np.sqrt(2. / (filter_size**2 * in_filters))
w = layers.weight_variable(
[filter_size, filter_size, in_filters, out_filters],
stddev=stddev,
name="weights")
out_ = tf.nn.conv2d(out_,
w, [1, 1, 1, 1],
padding="SAME",
data_format="NCHW")
out_ = out_ + layers.bias_variable([out_filters, 1, 1],
name='biases')
# Res Block 1
conv_out = layers.res_block(out_,
filter_size,
channel_multiplier=1,
depthwise_multiplier=2,
convolutions=2,
training=training,
activation=activation,
batch_norm=batch_norm,
data_format="NCHW")
pool_out = layers.max_pool(conv_out, pool_size, data_format="NCHW")
bottom_out = layers.dropout(pool_out, keep_prob)
side_out = layers.dropout(conv_out, keep_prob)
outputs.append(side_out)
# Build contracting path
for i in xrange(1, depth):
bottom_out, side_out = step_down('step_' + str(i),
bottom_out,
filter_size=filter_size,
res_blocks=res_blocks,
keep_prob=keep_prob,
training=training)
outputs.append(side_out)
# Bottom [conv -> conv]
with tf.variable_scope('step_' + str(depth)):
with tf.variable_scope("res_block_0"):
conv_out, tiled_input = layers.res_block(bottom_out,
filter_size,
channel_multiplier=2,
depthwise_multiplier=2,
convolutions=2,
training=training,
activation=activation,
batch_norm=batch_norm,
data_format="NCHW")
for i in xrange(1, res_blocks):
with tf.variable_scope("res_block_" + str(i)):
conv_out = layers.res_block(conv_out,
filter_size,
channel_multiplier=1,
depthwise_multiplier=2,
convolutions=2,
training=training,
activation=activation,
batch_norm=batch_norm,
data_format="NCHW")
conv_out = conv_out + tiled_input
current_tensor = layers.dropout(conv_out, keep_prob)
with tf.variable_scope('expanding'):
# Set initial parameter
outputs.reverse()
# Build expanding path
for i in xrange(depth):
current_tensor = step_up('step_' + str(depth + i + 1),
current_tensor,
outputs[i],
filter_size=filter_size,
res_blocks=res_blocks,
keep_prob=keep_prob,
training=training)
# Last layer is a 1x1 convolution to get the predictions
# We don't want an activation function for this one (softmax will be applied later), so we're doing it manually
in_filters = current_tensor.shape.as_list()[1]
stddev = np.sqrt(2. / in_filters)
with tf.variable_scope('classification'):
w = layers.weight_variable([1, 1, in_filters, out_channels],
stddev,
name='weights')
b = layers.bias_variable([out_channels, 1, 1], name='biases')
conv = tf.nn.conv2d(current_tensor,
w,
strides=[1, 1, 1, 1],
padding="SAME",
data_format="NCHW",
name='conv')
logits = conv + b
logits = tf.transpose(logits, perm=[0, 2, 3, 1])
return inputs, logits, ground_truth, keep_prob, training
