https://github.com/MarkMoHR/virtual_sketching
Tip revision: 958efe45a9120b9d467ba7701efba28c11e38f8f authored by Your Name on 21 August 2021, 08:13:15 UTC
Added bilibili links
Added bilibili links
Tip revision: 958efe4
model_common_train.py
import rnn
import tensorflow as tf
from subnet_tf_utils import generative_cnn_encoder, generative_cnn_encoder_deeper, generative_cnn_encoder_deeper13, \
generative_cnn_c3_encoder, generative_cnn_c3_encoder_deeper, generative_cnn_c3_encoder_deeper13, \
generative_cnn_c3_encoder_combine33, generative_cnn_c3_encoder_combine43, \
generative_cnn_c3_encoder_combine53, generative_cnn_c3_encoder_combineFC, \
generative_cnn_c3_encoder_deeper13_attn
from rasterization_utils.NeuralRenderer import NeuralRasterizorStep
from vgg_utils.VGG16 import vgg_net_slim
class VirtualSketchingModel(object):
def __init__(self, hps, gpu_mode=True, reuse=False):
"""Initializer for the model.
Args:
hps: a HParams object containing model hyperparameters
gpu_mode: a boolean that when True, uses GPU mode.
reuse: a boolean that when true, attemps to reuse variables.
"""
self.hps = hps
assert hps.model_mode in ['train', 'eval', 'eval_sample', 'sample']
# with tf.variable_scope('SCC', reuse=reuse):
if not gpu_mode:
with tf.device('/cpu:0'):
print('Model using cpu.')
self.build_model()
else:
print('-' * 100)
print('model_mode:', hps.model_mode)
print('Model using gpu.')
self.build_model()
def build_model(self):
"""Define model architecture."""
self.config_model()
initial_state = self.get_decoder_inputs()
self.initial_state = initial_state
self.initial_state_list = tf.split(self.initial_state, self.total_loop, axis=0)
total_loss_list = []
ras_loss_list = []
perc_relu_raw_list = []
perc_relu_norm_list = []
sn_loss_list = []
cursor_outside_loss_list = []
win_size_outside_loss_list = []
early_state_loss_list = []
tower_grads = []
pred_raster_imgs_list = []
pred_raster_imgs_rgb_list = []
for t_i in range(self.total_loop):
gpu_idx = t_i // self.hps.loop_per_gpu
gpu_i = self.hps.gpus[gpu_idx]
print(self.hps.model_mode, 'model, gpu:', gpu_i, ', loop:', t_i % self.hps.loop_per_gpu)
with tf.device('/gpu:%d' % gpu_i):
with tf.name_scope('GPU_%d' % gpu_i) as scope:
if t_i > 0:
tf.get_variable_scope().reuse_variables()
else:
total_loss_list.clear()
ras_loss_list.clear()
perc_relu_raw_list.clear()
perc_relu_norm_list.clear()
sn_loss_list.clear()
cursor_outside_loss_list.clear()
win_size_outside_loss_list.clear()
early_state_loss_list.clear()
tower_grads.clear()
pred_raster_imgs_list.clear()
pred_raster_imgs_rgb_list.clear()
split_input_photo = self.input_photo_list[t_i]
split_image_size = self.image_size[t_i]
split_init_cursor = self.init_cursor_list[t_i]
split_initial_state = self.initial_state_list[t_i]
if self.hps.input_channel == 1:
split_target_sketch = split_input_photo
else:
split_target_sketch = self.target_sketch_list[t_i]
## use pred as the prev points
other_params, pen_ras, final_state, pred_raster_images, pred_raster_images_rgb, \
pos_before_max_min, win_size_before_max_min \
= self.get_points_and_raster_image(split_initial_state, split_init_cursor, split_input_photo,
split_image_size)
# other_params: (N * max_seq_len, 6)
# pen_ras: (N * max_seq_len, 2), after softmax
# pos_before_max_min: (N, max_seq_len, 2), in image_size
# win_size_before_max_min: (N, max_seq_len, 1), in image_size
pred_raster_imgs = 1.0 - pred_raster_images # (N, image_size, image_size), [0.0-stroke, 1.0-BG]
pred_raster_imgs_rgb = 1.0 - pred_raster_images_rgb # (N, image_size, image_size, 3)
pred_raster_imgs_list.append(pred_raster_imgs)
pred_raster_imgs_rgb_list.append(pred_raster_imgs_rgb)
if not self.hps.use_softargmax:
pen_state_soft = pen_ras[:, 1:2] # (N * max_seq_len, 1)
else:
pen_state_soft = self.differentiable_argmax(pen_ras, self.hps.soft_beta) # (N * max_seq_len, 1)
pred_params = tf.concat([pen_state_soft, other_params], axis=1) # (N * max_seq_len, 7)
pred_params = tf.reshape(pred_params, shape=[-1, self.hps.max_seq_len, 7]) # (N, max_seq_len, 7)
# pred_params: (N, max_seq_len, 7)
if self.hps.model_mode == 'train' or self.hps.model_mode == 'eval':
raster_cost, sn_cost, cursor_outside_cost, winsize_outside_cost, \
early_pen_states_cost, \
perc_relu_loss_raw, perc_relu_loss_norm = \
self.build_losses(split_target_sketch, pred_raster_imgs, pred_params,
pos_before_max_min, win_size_before_max_min,
split_image_size)
# perc_relu_loss_raw, perc_relu_loss_norm: (n_layers)
ras_loss_list.append(raster_cost)
perc_relu_raw_list.append(perc_relu_loss_raw)
perc_relu_norm_list.append(perc_relu_loss_norm)
sn_loss_list.append(sn_cost)
cursor_outside_loss_list.append(cursor_outside_cost)
win_size_outside_loss_list.append(winsize_outside_cost)
early_state_loss_list.append(early_pen_states_cost)
if self.hps.model_mode == 'train':
total_cost_split, grads_and_vars_split = self.build_training_op_split(
raster_cost, sn_cost, cursor_outside_cost, winsize_outside_cost,
early_pen_states_cost)
total_loss_list.append(total_cost_split)
tower_grads.append(grads_and_vars_split)
self.raster_cost = tf.reduce_mean(tf.stack(ras_loss_list, axis=0))
self.perc_relu_losses_raw = tf.reduce_mean(tf.stack(perc_relu_raw_list, axis=0), axis=0) # (n_layers)
self.perc_relu_losses_norm = tf.reduce_mean(tf.stack(perc_relu_norm_list, axis=0), axis=0) # (n_layers)
self.stroke_num_cost = tf.reduce_mean(tf.stack(sn_loss_list, axis=0))
self.pos_outside_cost = tf.reduce_mean(tf.stack(cursor_outside_loss_list, axis=0))
self.win_size_outside_cost = tf.reduce_mean(tf.stack(win_size_outside_loss_list, axis=0))
self.early_pen_states_cost = tf.reduce_mean(tf.stack(early_state_loss_list, axis=0))
self.cost = tf.reduce_mean(tf.stack(total_loss_list, axis=0))
self.pred_raster_imgs = tf.concat(pred_raster_imgs_list, axis=0) # (N, image_size, image_size), [0.0-stroke, 1.0-BG]
self.pred_raster_imgs_rgb = tf.concat(pred_raster_imgs_rgb_list, axis=0) # (N, image_size, image_size, 3)
if self.hps.model_mode == 'train':
self.build_training_op(tower_grads)
def config_model(self):
if self.hps.model_mode == 'train':
self.global_step = tf.Variable(0, name='global_step', trainable=False)
if self.hps.dec_model == 'lstm':
dec_cell_fn = rnn.LSTMCell
elif self.hps.dec_model == 'layer_norm':
dec_cell_fn = rnn.LayerNormLSTMCell
elif self.hps.dec_model == 'hyper':
dec_cell_fn = rnn.HyperLSTMCell
else:
assert False, 'please choose a respectable cell'
use_recurrent_dropout = self.hps.use_recurrent_dropout
use_input_dropout = self.hps.use_input_dropout
use_output_dropout = self.hps.use_output_dropout
dec_cell = dec_cell_fn(
self.hps.dec_rnn_size,
use_recurrent_dropout=use_recurrent_dropout,
dropout_keep_prob=self.hps.recurrent_dropout_prob)
# dropout:
# print('Input dropout mode = %s.' % use_input_dropout)
# print('Output dropout mode = %s.' % use_output_dropout)
# print('Recurrent dropout mode = %s.' % use_recurrent_dropout)
if use_input_dropout:
print('Dropout to input w/ keep_prob = %4.4f.' % self.hps.input_dropout_prob)
dec_cell = tf.contrib.rnn.DropoutWrapper(
dec_cell, input_keep_prob=self.hps.input_dropout_prob)
if use_output_dropout:
print('Dropout to output w/ keep_prob = %4.4f.' % self.hps.output_dropout_prob)
dec_cell = tf.contrib.rnn.DropoutWrapper(
dec_cell, output_keep_prob=self.hps.output_dropout_prob)
self.dec_cell = dec_cell
self.total_loop = len(self.hps.gpus) * self.hps.loop_per_gpu
self.init_cursor = tf.placeholder(
dtype=tf.float32,
shape=[self.hps.batch_size, 1, 2]) # (N, 1, 2), in size [0.0, 1.0)
self.init_width = tf.placeholder(
dtype=tf.float32,
shape=[1]) # (1), in [0.0, 1.0]
self.image_size = tf.placeholder(dtype=tf.int32, shape=(self.total_loop)) # ()
self.init_cursor_list = tf.split(self.init_cursor, self.total_loop, axis=0)
self.input_photo_list = []
for loop_i in range(self.total_loop):
input_photo_i = tf.placeholder(dtype=tf.float32, shape=[None, None, None, self.hps.input_channel]) # [0.0-stroke, 1.0-BG]
self.input_photo_list.append(input_photo_i)
if self.hps.input_channel == 3:
self.target_sketch_list = []
for loop_i in range(self.total_loop):
target_sketch_i = tf.placeholder(dtype=tf.float32, shape=[None, None, None, 1]) # [0.0-stroke, 1.0-BG]
self.target_sketch_list.append(target_sketch_i)
if self.hps.model_mode == 'train' or self.hps.model_mode == 'eval':
self.stroke_num_loss_weight = tf.Variable(0.0, trainable=False)
self.early_pen_loss_start_idx = tf.Variable(0, dtype=tf.int32, trainable=False)
self.early_pen_loss_end_idx = tf.Variable(0, dtype=tf.int32, trainable=False)
if self.hps.model_mode == 'train':
self.perc_loss_mean_list = []
for loop_i in range(len(self.hps.perc_loss_layers)):
relu_loss_mean = tf.Variable(0.0, trainable=False)
self.perc_loss_mean_list.append(relu_loss_mean)
self.last_step_num = tf.Variable(0.0, trainable=False)
with tf.variable_scope('train_op', reuse=tf.AUTO_REUSE):
self.lr = tf.Variable(self.hps.learning_rate, trainable=False)
self.optimizer = tf.train.AdamOptimizer(self.lr)
###########################
def normalize_image_m1to1(self, in_img_0to1):
norm_img_m1to1 = tf.multiply(in_img_0to1, 2.0)
norm_img_m1to1 = tf.subtract(norm_img_m1to1, 1.0)
return norm_img_m1to1
def add_coords(self, input_tensor):
batch_size_tensor = tf.shape(input_tensor)[0] # get N size
xx_ones = tf.ones([batch_size_tensor, self.hps.raster_size], dtype=tf.int32) # e.g. (N, raster_size)
xx_ones = tf.expand_dims(xx_ones, -1) # e.g. (N, raster_size, 1)
xx_range = tf.tile(tf.expand_dims(tf.range(self.hps.raster_size), 0),
[batch_size_tensor, 1]) # e.g. (N, raster_size)
xx_range = tf.expand_dims(xx_range, 1) # e.g. (N, 1, raster_size)
xx_channel = tf.matmul(xx_ones, xx_range) # e.g. (N, raster_size, raster_size)
xx_channel = tf.expand_dims(xx_channel, -1) # e.g. (N, raster_size, raster_size, 1)
yy_ones = tf.ones([batch_size_tensor, self.hps.raster_size], dtype=tf.int32) # e.g. (N, raster_size)
yy_ones = tf.expand_dims(yy_ones, 1) # e.g. (N, 1, raster_size)
yy_range = tf.tile(tf.expand_dims(tf.range(self.hps.raster_size), 0),
[batch_size_tensor, 1]) # (N, raster_size)
yy_range = tf.expand_dims(yy_range, -1) # e.g. (N, raster_size, 1)
yy_channel = tf.matmul(yy_range, yy_ones) # e.g. (N, raster_size, raster_size)
yy_channel = tf.expand_dims(yy_channel, -1) # e.g. (N, raster_size, raster_size, 1)
xx_channel = tf.cast(xx_channel, 'float32') / (self.hps.raster_size - 1)
yy_channel = tf.cast(yy_channel, 'float32') / (self.hps.raster_size - 1)
# xx_channel = xx_channel * 2 - 1 # [-1, 1]
# yy_channel = yy_channel * 2 - 1
ret = tf.concat([
input_tensor,
xx_channel,
yy_channel,
], axis=-1) # e.g. (N, raster_size, raster_size, 4)
return ret
def build_combined_encoder(self, patch_canvas, patch_photo, entire_canvas, entire_photo, cursor_pos,
image_size, window_size):
"""
:param patch_canvas: (N, raster_size, raster_size, 1), [-1.0-stroke, 1.0-BG]
:param patch_photo: (N, raster_size, raster_size, 1/3), [-1.0-stroke, 1.0-BG]
:param entire_canvas: (N, image_size, image_size, 1), [0.0-stroke, 1.0-BG]
:param entire_photo: (N, image_size, image_size, 1/3), [0.0-stroke, 1.0-BG]
:param cursor_pos: (N, 1, 2), in size [0.0, 1.0)
:param window_size: (N, 1, 1), float, in large size
:return:
"""
if self.hps.resize_method == 'BILINEAR':
resize_method = tf.image.ResizeMethod.BILINEAR
elif self.hps.resize_method == 'NEAREST_NEIGHBOR':
resize_method = tf.image.ResizeMethod.NEAREST_NEIGHBOR
elif self.hps.resize_method == 'BICUBIC':
resize_method = tf.image.ResizeMethod.BICUBIC
elif self.hps.resize_method == 'AREA':
resize_method = tf.image.ResizeMethod.AREA
else:
raise Exception('unknown resize_method', self.hps.resize_method)
patch_photo = tf.stop_gradient(patch_photo)
patch_canvas = tf.stop_gradient(patch_canvas)
cursor_pos = tf.stop_gradient(cursor_pos)
window_size = tf.stop_gradient(window_size)
entire_photo_small = tf.stop_gradient(tf.image.resize_images(entire_photo,
(self.hps.raster_size, self.hps.raster_size),
method=resize_method))
entire_canvas_small = tf.stop_gradient(tf.image.resize_images(entire_canvas,
(self.hps.raster_size, self.hps.raster_size),
method=resize_method))
entire_photo_small = self.normalize_image_m1to1(entire_photo_small) # [-1.0-stroke, 1.0-BG]
entire_canvas_small = self.normalize_image_m1to1(entire_canvas_small) # [-1.0-stroke, 1.0-BG]
if self.hps.encode_cursor_type == 'value':
cursor_pos_norm = tf.expand_dims(cursor_pos, axis=1) # (N, 1, 1, 2)
cursor_pos_norm = tf.tile(cursor_pos_norm, [1, self.hps.raster_size, self.hps.raster_size, 1])
cursor_info = cursor_pos_norm
else:
raise Exception('Unknown encode_cursor_type', self.hps.encode_cursor_type)
batch_input_combined = tf.concat([patch_photo, patch_canvas, entire_photo_small, entire_canvas_small, cursor_info],
axis=-1) # [N, raster_size, raster_size, 6/10]
batch_input_local = tf.concat([patch_photo, patch_canvas], axis=-1) # [N, raster_size, raster_size, 2/4]
batch_input_global = tf.concat([entire_photo_small, entire_canvas_small, cursor_info],
axis=-1) # [N, raster_size, raster_size, 4/6]
if self.hps.model_mode == 'train':
is_training = True
dropout_keep_prob = self.hps.pix_drop_kp
else:
is_training = False
dropout_keep_prob = 1.0
if self.hps.add_coordconv:
batch_input_combined = self.add_coords(batch_input_combined) # (N, in_H, in_W, in_dim + 2)
batch_input_local = self.add_coords(batch_input_local) # (N, in_H, in_W, in_dim + 2)
batch_input_global = self.add_coords(batch_input_global) # (N, in_H, in_W, in_dim + 2)
if 'combine' in self.hps.encoder_type:
if self.hps.encoder_type == 'combine33':
image_embedding, _ = generative_cnn_c3_encoder_combine33(batch_input_local, batch_input_global,
is_training, dropout_keep_prob) # (N, 128)
elif self.hps.encoder_type == 'combine43':
image_embedding, _ = generative_cnn_c3_encoder_combine43(batch_input_local, batch_input_global,
is_training, dropout_keep_prob) # (N, 128)
elif self.hps.encoder_type == 'combine53':
image_embedding, _ = generative_cnn_c3_encoder_combine53(batch_input_local, batch_input_global,
is_training, dropout_keep_prob) # (N, 128)
elif self.hps.encoder_type == 'combineFC':
image_embedding, _ = generative_cnn_c3_encoder_combineFC(batch_input_local, batch_input_global,
is_training, dropout_keep_prob) # (N, 256)
else:
raise Exception('Unknown encoder_type', self.hps.encoder_type)
else:
with tf.variable_scope('Combined_Encoder', reuse=tf.AUTO_REUSE):
if self.hps.encoder_type == 'conv10':
image_embedding, _ = generative_cnn_encoder(batch_input_combined, is_training, dropout_keep_prob) # (N, 128)
elif self.hps.encoder_type == 'conv10_deep':
image_embedding, _ = generative_cnn_encoder_deeper(batch_input_combined, is_training, dropout_keep_prob) # (N, 512)
elif self.hps.encoder_type == 'conv13':
image_embedding, _ = generative_cnn_encoder_deeper13(batch_input_combined, is_training, dropout_keep_prob) # (N, 128)
elif self.hps.encoder_type == 'conv10_c3':
image_embedding, _ = generative_cnn_c3_encoder(batch_input_combined, is_training, dropout_keep_prob) # (N, 128)
elif self.hps.encoder_type == 'conv10_deep_c3':
image_embedding, _ = generative_cnn_c3_encoder_deeper(batch_input_combined, is_training, dropout_keep_prob) # (N, 512)
elif self.hps.encoder_type == 'conv13_c3':
image_embedding, _ = generative_cnn_c3_encoder_deeper13(batch_input_combined, is_training, dropout_keep_prob) # (N, 128)
elif self.hps.encoder_type == 'conv13_c3_attn':
image_embedding, _ = generative_cnn_c3_encoder_deeper13_attn(batch_input_combined, is_training, dropout_keep_prob) # (N, 128)
else:
raise Exception('Unknown encoder_type', self.hps.encoder_type)
return image_embedding
def build_seq_decoder(self, dec_cell, actual_input_x, initial_state):
rnn_output, last_state = self.rnn_decoder(dec_cell, initial_state, actual_input_x)
rnn_output_flat = tf.reshape(rnn_output, [-1, self.hps.dec_rnn_size])
pen_n_out = 2
params_n_out = 6
with tf.variable_scope('DEC_RNN_out_pen', reuse=tf.AUTO_REUSE):
output_w_pen = tf.get_variable('output_w', [self.hps.dec_rnn_size, pen_n_out])
output_b_pen = tf.get_variable('output_b', [pen_n_out], initializer=tf.constant_initializer(0.0))
output_pen = tf.nn.xw_plus_b(rnn_output_flat, output_w_pen, output_b_pen) # (N, pen_n_out)
with tf.variable_scope('DEC_RNN_out_params', reuse=tf.AUTO_REUSE):
output_w_params = tf.get_variable('output_w', [self.hps.dec_rnn_size, params_n_out])
output_b_params = tf.get_variable('output_b', [params_n_out], initializer=tf.constant_initializer(0.0))
output_params = tf.nn.xw_plus_b(rnn_output_flat, output_w_params, output_b_params) # (N, params_n_out)
output = tf.concat([output_pen, output_params], axis=1) # (N, n_out)
return output, last_state
def get_mixture_coef(self, outputs):
z = outputs
z_pen_logits = z[:, 0:2] # (N, 2), pen states
z_other_params_logits = z[:, 2:] # (N, 6)
z_pen = tf.nn.softmax(z_pen_logits) # (N, 2)
if self.hps.position_format == 'abs':
x1y1 = tf.nn.sigmoid(z_other_params_logits[:, 0:2]) # (N, 2)
x2y2 = tf.tanh(z_other_params_logits[:, 2:4]) # (N, 2)
widths = tf.nn.sigmoid(z_other_params_logits[:, 4:5]) # (N, 1)
widths = tf.add(tf.multiply(widths, 1.0 - self.hps.min_width), self.hps.min_width)
scaling = tf.nn.sigmoid(z_other_params_logits[:, 5:6]) * self.hps.max_scaling # (N, 1), [0.0, max_scaling]
# scaling = tf.add(tf.multiply(scaling, (self.hps.max_scaling - self.hps.min_scaling) / self.hps.max_scaling),
# self.hps.min_scaling)
z_other_params = tf.concat([x1y1, x2y2, widths, scaling], axis=-1) # (N, 6)
else: # "rel"
raise Exception('Unknown position_format', self.hps.position_format)
r = [z_other_params, z_pen]
return r
###########################
def get_decoder_inputs(self):
initial_state = self.dec_cell.zero_state(batch_size=self.hps.batch_size, dtype=tf.float32)
return initial_state
def rnn_decoder(self, dec_cell, initial_state, actual_input_x):
with tf.variable_scope("RNN_DEC", reuse=tf.AUTO_REUSE):
output, last_state = tf.nn.dynamic_rnn(
dec_cell,
actual_input_x,
initial_state=initial_state,
time_major=False,
swap_memory=True,
dtype=tf.float32)
return output, last_state
###########################
def image_padding(self, ori_image, window_size, pad_value):
"""
Pad with (bg)
:param ori_image:
:return:
"""
paddings = [[0, 0],
[window_size // 2, window_size // 2],
[window_size // 2, window_size // 2],
[0, 0]]
pad_img = tf.pad(ori_image, paddings=paddings, mode='CONSTANT', constant_values=pad_value) # (N, H_p, W_p, k)
return pad_img
def image_cropping_fn(self, fn_inputs):
"""
crop the patch
:return:
"""
index_offset = self.hps.input_channel - 1
input_image = fn_inputs[:, :, 0:2 + index_offset] # (image_size, image_size, 2), [0.0-BG, 1.0-stroke]
cursor_pos = fn_inputs[0, 0, 2 + index_offset:4 + index_offset] # (2), in [0.0, 1.0)
image_size = fn_inputs[0, 0, 4 + index_offset] # (), float32
window_size = tf.cast(fn_inputs[0, 0, 5 + index_offset], tf.int32) # ()
input_img_reshape = tf.expand_dims(input_image, axis=0)
pad_img = self.image_padding(input_img_reshape, window_size, pad_value=0.0)
cursor_pos = tf.cast(tf.round(tf.multiply(cursor_pos, image_size)), dtype=tf.int32)
x0, x1 = cursor_pos[0], cursor_pos[0] + window_size # ()
y0, y1 = cursor_pos[1], cursor_pos[1] + window_size # ()
patch_image = pad_img[:, y0:y1, x0:x1, :] # (1, window_size, window_size, 2/4)
# resize to raster_size
patch_image_scaled = tf.image.resize_images(patch_image, (self.hps.raster_size, self.hps.raster_size),
method=tf.image.ResizeMethod.AREA)
patch_image_scaled = tf.squeeze(patch_image_scaled, axis=0)
# patch_canvas_scaled: (raster_size, raster_size, 2/4), [0.0-BG, 1.0-stroke]
return patch_image_scaled
def image_cropping(self, cursor_position, input_img, image_size, window_sizes):
"""
:param cursor_position: (N, 1, 2), float type, in size [0.0, 1.0)
:param input_img: (N, image_size, image_size, 2/4), [0.0-BG, 1.0-stroke]
:param window_sizes: (N, 1, 1), float32, with grad
"""
input_img_ = input_img
window_sizes_non_grad = tf.stop_gradient(tf.round(window_sizes)) # (N, 1, 1), no grad
cursor_position_ = tf.reshape(cursor_position, (-1, 1, 1, 2)) # (N, 1, 1, 2)
cursor_position_ = tf.tile(cursor_position_, [1, image_size, image_size, 1]) # (N, image_size, image_size, 2)
image_size_ = tf.reshape(tf.cast(image_size, tf.float32), (1, 1, 1, 1)) # (1, 1, 1, 1)
image_size_ = tf.tile(image_size_, [self.hps.batch_size // self.total_loop, image_size, image_size, 1])
window_sizes_ = tf.reshape(window_sizes_non_grad, (-1, 1, 1, 1)) # (N, 1, 1, 1)
window_sizes_ = tf.tile(window_sizes_, [1, image_size, image_size, 1]) # (N, image_size, image_size, 1)
fn_inputs = tf.concat([input_img_, cursor_position_, image_size_, window_sizes_],
axis=-1) # (N, image_size, image_size, 2/4 + 4)
curr_patch_imgs = tf.map_fn(self.image_cropping_fn, fn_inputs, parallel_iterations=32) # (N, raster_size, raster_size, -)
return curr_patch_imgs
def image_cropping_v3(self, cursor_position, input_img, image_size, window_sizes):
"""
:param cursor_position: (N, 1, 2), float type, in size [0.0, 1.0)
:param input_img: (N, image_size, image_size, k), [0.0-BG, 1.0-stroke]
:param window_sizes: (N, 1, 1), float32, with grad
"""
window_sizes_non_grad = tf.stop_gradient(window_sizes) # (N, 1, 1), no grad
cursor_pos = tf.multiply(cursor_position, tf.cast(image_size, tf.float32))
cursor_x, cursor_y = tf.split(cursor_pos, 2, axis=-1) # (N, 1, 1)
y1 = cursor_y - (window_sizes_non_grad - 1.0) / 2
x1 = cursor_x - (window_sizes_non_grad - 1.0) / 2
y2 = y1 + (window_sizes_non_grad - 1.0)
x2 = x1 + (window_sizes_non_grad - 1.0)
boxes = tf.concat([y1, x1, y2, x2], axis=-1) # (N, 1, 4)
boxes = tf.squeeze(boxes, axis=1) # (N, 4)
boxes = boxes / tf.cast(image_size - 1, tf.float32)
box_ind = tf.ones_like(cursor_x)[:, 0, 0] # (N)
box_ind = tf.cast(box_ind, dtype=tf.int32)
box_ind = tf.cumsum(box_ind) - 1
curr_patch_imgs = tf.image.crop_and_resize(input_img, boxes, box_ind,
crop_size=[self.hps.raster_size, self.hps.raster_size])
# (N, raster_size, raster_size, k), [0.0-BG, 1.0-stroke]
return curr_patch_imgs
def get_pixel_value(self, img, x, y):
"""
Utility function to get pixel value for coordinate vectors x and y from a 4D tensor image.
Input
-----
- img: tensor of shape (B, H, W, C)
- x: flattened tensor of shape (B, H', W')
- y: flattened tensor of shape (B, H', W')
Returns
-------
- output: tensor of shape (B, H', W', C)
"""
shape = tf.shape(x)
batch_size = shape[0]
height = shape[1]
width = shape[2]
batch_idx = tf.range(0, batch_size)
batch_idx = tf.reshape(batch_idx, (batch_size, 1, 1))
b = tf.tile(batch_idx, (1, height, width))
indices = tf.stack([b, y, x], 3)
return tf.gather_nd(img, indices)
def image_pasting_nondiff_single(self, fn_inputs):
patch_image = fn_inputs[:, :, 0:1] # (raster_size, raster_size, 1), [0.0-BG, 1.0-stroke]
cursor_pos = fn_inputs[0, 0, 1:3] # (2), in large size
image_size = tf.cast(fn_inputs[0, 0, 3], tf.int32) # ()
window_size = tf.cast(fn_inputs[0, 0, 4], tf.int32) # ()
patch_image_scaled = tf.expand_dims(patch_image, axis=0) # (1, raster_size, raster_size, 1)
patch_image_scaled = tf.image.resize_images(patch_image_scaled, (window_size, window_size),
method=tf.image.ResizeMethod.BILINEAR)
patch_image_scaled = tf.squeeze(patch_image_scaled, axis=0)
# patch_canvas_scaled: (window_size, window_size, 1)
cursor_pos = tf.cast(tf.round(cursor_pos), dtype=tf.int32) # (2)
cursor_x, cursor_y = cursor_pos[0], cursor_pos[1]
pad_up = cursor_y
pad_down = image_size - cursor_y
pad_left = cursor_x
pad_right = image_size - cursor_x
paddings = [[pad_up, pad_down],
[pad_left, pad_right],
[0, 0]]
pad_img = tf.pad(patch_image_scaled, paddings=paddings, mode='CONSTANT',
constant_values=0.0) # (H_p, W_p, 1), [0.0-BG, 1.0-stroke]
crop_start = window_size // 2
pasted_image = pad_img[crop_start: crop_start + image_size, crop_start: crop_start + image_size, :]
return pasted_image
def image_pasting_diff_single(self, fn_inputs):
patch_canvas = fn_inputs[:, :, 0:1] # (raster_size, raster_size, 1), [0.0-BG, 1.0-stroke]
cursor_pos = fn_inputs[0, 0, 1:3] # (2), in large size
image_size = tf.cast(fn_inputs[0, 0, 3], tf.int32) # ()
window_size = tf.cast(fn_inputs[0, 0, 4], tf.int32) # ()
cursor_x, cursor_y = cursor_pos[0], cursor_pos[1]
patch_canvas_scaled = tf.expand_dims(patch_canvas, axis=0) # (1, raster_size, raster_size, 1)
patch_canvas_scaled = tf.image.resize_images(patch_canvas_scaled, (window_size, window_size),
method=tf.image.ResizeMethod.BILINEAR)
# patch_canvas_scaled: (1, window_size, window_size, 1)
valid_canvas = self.image_pasting_diff_batch(patch_canvas_scaled,
tf.expand_dims(tf.expand_dims(cursor_pos, axis=0), axis=0),
window_size)
valid_canvas = tf.squeeze(valid_canvas, axis=0)
# (window_size + 1, window_size + 1, 1)
pad_up = tf.cast(tf.floor(cursor_y), tf.int32)
pad_down = image_size - 1 - tf.cast(tf.floor(cursor_y), tf.int32)
pad_left = tf.cast(tf.floor(cursor_x), tf.int32)
pad_right = image_size - 1 - tf.cast(tf.floor(cursor_x), tf.int32)
paddings = [[pad_up, pad_down],
[pad_left, pad_right],
[0, 0]]
pad_img = tf.pad(valid_canvas, paddings=paddings, mode='CONSTANT',
constant_values=0.0) # (H_p, W_p, 1), [0.0-BG, 1.0-stroke]
crop_start = window_size // 2
pasted_image = pad_img[crop_start: crop_start + image_size, crop_start: crop_start + image_size, :]
return pasted_image
def image_pasting_diff_single_v3(self, fn_inputs):
patch_canvas = fn_inputs[:, :, 0:1] # (raster_size, raster_size, 1), [0.0-BG, 1.0-stroke]
cursor_pos_a = fn_inputs[0, 0, 1:3] # (2), float32, in large size
image_size_a = tf.cast(fn_inputs[0, 0, 3], tf.int32) # ()
window_size_a = fn_inputs[0, 0, 4] # (), float32, with grad
raster_size_a = float(self.hps.raster_size)
padding_size = tf.cast(tf.ceil(window_size_a / 2.0), tf.int32)
x1y1_a = cursor_pos_a - window_size_a / 2.0 # (2), float32
x2y2_a = cursor_pos_a + window_size_a / 2.0 # (2), float32
x1y1_a_floor = tf.floor(x1y1_a) # (2)
x2y2_a_ceil = tf.ceil(x2y2_a) # (2)
cursor_pos_b_oricoord = (x1y1_a_floor + x2y2_a_ceil) / 2.0 # (2)
cursor_pos_b = (cursor_pos_b_oricoord - x1y1_a) / window_size_a * raster_size_a # (2)
raster_size_b = (x2y2_a_ceil - x1y1_a_floor) # (x, y)
image_size_b = raster_size_a
window_size_b = raster_size_a * (raster_size_b / window_size_a) # (x, y)
cursor_b_x, cursor_b_y = tf.split(cursor_pos_b, 2, axis=-1) # (1)
y1_b = cursor_b_y - (window_size_b[1] - 1.) / 2.
x1_b = cursor_b_x - (window_size_b[0] - 1.) / 2.
y2_b = y1_b + (window_size_b[1] - 1.)
x2_b = x1_b + (window_size_b[0] - 1.)
boxes_b = tf.concat([y1_b, x1_b, y2_b, x2_b], axis=-1) # (4)
boxes_b = boxes_b / tf.cast(image_size_b - 1, tf.float32) # with grad to window_size_a
box_ind_b = tf.ones((1), dtype=tf.int32) # (1)
box_ind_b = tf.cumsum(box_ind_b) - 1
patch_canvas = tf.expand_dims(patch_canvas, axis=0) # (1, raster_size, raster_size, 1), [0.0-BG, 1.0-stroke]
boxes_b = tf.expand_dims(boxes_b, axis=0) # (1, 4)
valid_canvas = tf.image.crop_and_resize(patch_canvas, boxes_b, box_ind_b,
crop_size=[raster_size_b[1], raster_size_b[0]])
valid_canvas = valid_canvas[0] # (raster_size_b, raster_size_b, 1)
pad_up = tf.cast(x1y1_a_floor[1], tf.int32) + padding_size
pad_down = image_size_a + padding_size - tf.cast(x2y2_a_ceil[1], tf.int32)
pad_left = tf.cast(x1y1_a_floor[0], tf.int32) + padding_size
pad_right = image_size_a + padding_size - tf.cast(x2y2_a_ceil[0], tf.int32)
paddings = [[pad_up, pad_down],
[pad_left, pad_right],
[0, 0]]
pad_img = tf.pad(valid_canvas, paddings=paddings, mode='CONSTANT',
constant_values=0.0) # (H_p, W_p, 1), [0.0-BG, 1.0-stroke]
pasted_image = pad_img[padding_size: padding_size + image_size_a, padding_size: padding_size + image_size_a, :]
return pasted_image
def image_pasting_diff_batch(self, patch_image, cursor_position, window_size):
"""
:param patch_img: (N, window_size, window_size, 1), [0.0-BG, 1.0-stroke]
:param cursor_position: (N, 1, 2), in large size
:return:
"""
paddings1 = [[0, 0],
[1, 1],
[1, 1],
[0, 0]]
patch_image_pad1 = tf.pad(patch_image, paddings=paddings1, mode='CONSTANT',
constant_values=0.0) # (N, window_size+2, window_size+2, 1), [0.0-BG, 1.0-stroke]
cursor_x, cursor_y = cursor_position[:, :, 0:1], cursor_position[:, :, 1:2] # (N, 1, 1)
cursor_x_f, cursor_y_f = tf.floor(cursor_x), tf.floor(cursor_y)
patch_x, patch_y = 1.0 - (cursor_x - cursor_x_f), 1.0 - (cursor_y - cursor_y_f) # (N, 1, 1)
x_ones = tf.ones_like(patch_x, dtype=tf.float32) # (N, 1, 1)
x_ones = tf.tile(x_ones, [1, 1, window_size]) # (N, 1, window_size)
patch_x = tf.concat([patch_x, x_ones], axis=-1) # (N, 1, window_size + 1)
patch_x = tf.tile(patch_x, [1, window_size + 1, 1]) # (N, window_size + 1, window_size + 1)
patch_x = tf.cumsum(patch_x, axis=-1) # (N, window_size + 1, window_size + 1)
patch_x0 = tf.cast(tf.floor(patch_x), tf.int32) # (N, window_size + 1, window_size + 1)
patch_x1 = patch_x0 + 1 # (N, window_size + 1, window_size + 1)
y_ones = tf.ones_like(patch_y, dtype=tf.float32) # (N, 1, 1)
y_ones = tf.tile(y_ones, [1, window_size, 1]) # (N, window_size, 1)
patch_y = tf.concat([patch_y, y_ones], axis=1) # (N, window_size + 1, 1)
patch_y = tf.tile(patch_y, [1, 1, window_size + 1]) # (N, window_size + 1, window_size + 1)
patch_y = tf.cumsum(patch_y, axis=1) # (N, window_size + 1, window_size + 1)
patch_y0 = tf.cast(tf.floor(patch_y), tf.int32) # (N, window_size + 1, window_size + 1)
patch_y1 = patch_y0 + 1 # (N, window_size + 1, window_size + 1)
# get pixel value at corner coords
valid_canvas_patch_a = self.get_pixel_value(patch_image_pad1, patch_x0, patch_y0)
valid_canvas_patch_b = self.get_pixel_value(patch_image_pad1, patch_x0, patch_y1)
valid_canvas_patch_c = self.get_pixel_value(patch_image_pad1, patch_x1, patch_y0)
valid_canvas_patch_d = self.get_pixel_value(patch_image_pad1, patch_x1, patch_y1)
# (N, window_size + 1, window_size + 1, 1)
patch_x0 = tf.cast(patch_x0, tf.float32)
patch_x1 = tf.cast(patch_x1, tf.float32)
patch_y0 = tf.cast(patch_y0, tf.float32)
patch_y1 = tf.cast(patch_y1, tf.float32)
# calculate deltas
wa = (patch_x1 - patch_x) * (patch_y1 - patch_y)
wb = (patch_x1 - patch_x) * (patch_y - patch_y0)
wc = (patch_x - patch_x0) * (patch_y1 - patch_y)
wd = (patch_x - patch_x0) * (patch_y - patch_y0)
# (N, window_size + 1, window_size + 1)
# add dimension for addition
wa = tf.expand_dims(wa, axis=3)
wb = tf.expand_dims(wb, axis=3)
wc = tf.expand_dims(wc, axis=3)
wd = tf.expand_dims(wd, axis=3)
# (N, window_size + 1, window_size + 1, 1)
# compute output
valid_canvas_patch_ = tf.add_n([wa * valid_canvas_patch_a,
wb * valid_canvas_patch_b,
wc * valid_canvas_patch_c,
wd * valid_canvas_patch_d]) # (N, window_size + 1, window_size + 1, 1)
return valid_canvas_patch_
def image_pasting(self, cursor_position_norm, patch_img, image_size, window_sizes, is_differentiable=False):
"""
paste the patch_img to padded size based on cursor_position
:param cursor_position_norm: (N, 1, 2), float type, in size [0.0, 1.0)
:param patch_img: (N, raster_size, raster_size), [0.0-BG, 1.0-stroke]
:param window_sizes: (N, 1, 1), float32, with grad
:return:
"""
cursor_position = tf.multiply(cursor_position_norm, tf.cast(image_size, tf.float32)) # in large size
window_sizes_r = tf.round(window_sizes) # (N, 1, 1), no grad
patch_img_ = tf.expand_dims(patch_img, axis=-1) # (N, raster_size, raster_size, 1)
cursor_position_step = tf.reshape(cursor_position, (-1, 1, 1, 2)) # (N, 1, 1, 2)
cursor_position_step = tf.tile(cursor_position_step, [1, self.hps.raster_size, self.hps.raster_size,
1]) # (N, raster_size, raster_size, 2)
image_size_tile = tf.reshape(tf.cast(image_size, tf.float32), (1, 1, 1, 1)) # (N, 1, 1, 1)
image_size_tile = tf.tile(image_size_tile, [self.hps.batch_size // self.total_loop, self.hps.raster_size,
self.hps.raster_size, 1])
window_sizes_tile = tf.reshape(window_sizes_r, (-1, 1, 1, 1)) # (N, 1, 1, 1)
window_sizes_tile = tf.tile(window_sizes_tile, [1, self.hps.raster_size, self.hps.raster_size, 1])
pasting_inputs = tf.concat([patch_img_, cursor_position_step, image_size_tile, window_sizes_tile],
axis=-1) # (N, raster_size, raster_size, 5)
if is_differentiable:
curr_paste_imgs = tf.map_fn(self.image_pasting_diff_single, pasting_inputs,
parallel_iterations=32) # (N, image_size, image_size, 1)
else:
curr_paste_imgs = tf.map_fn(self.image_pasting_nondiff_single, pasting_inputs,
parallel_iterations=32) # (N, image_size, image_size, 1)
curr_paste_imgs = tf.squeeze(curr_paste_imgs, axis=-1) # (N, image_size, image_size)
return curr_paste_imgs
def image_pasting_v3(self, cursor_position_norm, patch_img, image_size, window_sizes, is_differentiable=False):
"""
paste the patch_img to padded size based on cursor_position
:param cursor_position_norm: (N, 1, 2), float type, in size [0.0, 1.0)
:param patch_img: (N, raster_size, raster_size), [0.0-BG, 1.0-stroke]
:param window_sizes: (N, 1, 1), float32, with grad
:return:
"""
cursor_position = tf.multiply(cursor_position_norm, tf.cast(image_size, tf.float32)) # in large size
if is_differentiable:
patch_img_ = tf.expand_dims(patch_img, axis=-1) # (N, raster_size, raster_size, 1)
cursor_position_step = tf.reshape(cursor_position, (-1, 1, 1, 2)) # (N, 1, 1, 2)
cursor_position_step = tf.tile(cursor_position_step, [1, self.hps.raster_size, self.hps.raster_size,
1]) # (N, raster_size, raster_size, 2)
image_size_tile = tf.reshape(tf.cast(image_size, tf.float32), (1, 1, 1, 1)) # (N, 1, 1, 1)
image_size_tile = tf.tile(image_size_tile, [self.hps.batch_size // self.total_loop, self.hps.raster_size,
self.hps.raster_size, 1])
window_sizes_tile = tf.reshape(window_sizes, (-1, 1, 1, 1)) # (N, 1, 1, 1)
window_sizes_tile = tf.tile(window_sizes_tile, [1, self.hps.raster_size, self.hps.raster_size, 1])
pasting_inputs = tf.concat([patch_img_, cursor_position_step, image_size_tile, window_sizes_tile],
axis=-1) # (N, raster_size, raster_size, 5)
curr_paste_imgs = tf.map_fn(self.image_pasting_diff_single_v3, pasting_inputs,
parallel_iterations=32) # (N, image_size, image_size, 1)
else:
raise Exception('Unfinished...')
curr_paste_imgs = tf.squeeze(curr_paste_imgs, axis=-1) # (N, image_size, image_size)
return curr_paste_imgs
def get_points_and_raster_image(self, initial_state, init_cursor, input_photo, image_size):
## generate the other_params and pen_ras and raster image for raster loss
prev_state = initial_state # (N, dec_rnn_size * 3)
prev_width = self.init_width # (1)
prev_width = tf.expand_dims(tf.expand_dims(prev_width, axis=0), axis=0) # (1, 1, 1)
prev_width = tf.tile(prev_width, [self.hps.batch_size // self.total_loop, 1, 1]) # (N, 1, 1)
prev_scaling = tf.ones((self.hps.batch_size // self.total_loop, 1, 1)) # (N, 1, 1)
prev_window_size = tf.ones((self.hps.batch_size // self.total_loop, 1, 1),
dtype=tf.float32) * float(self.hps.raster_size) # (N, 1, 1)
cursor_position_temp = init_cursor
self.cursor_position = cursor_position_temp # (N, 1, 2), in size [0.0, 1.0)
cursor_position_loop = self.cursor_position
other_params_list = []
pen_ras_list = []
pos_before_max_min_list = []
win_size_before_max_min_list = []
curr_canvas_soft = tf.zeros_like(input_photo[:, :, :, 0]) # (N, image_size, image_size), [0.0-BG, 1.0-stroke]
curr_canvas_soft_rgb = tf.tile(tf.zeros_like(input_photo[:, :, :, 0:1]), [1, 1, 1, 3]) # (N, image_size, image_size, 3), [0.0-BG, 1.0-stroke]
curr_canvas_hard = tf.zeros_like(curr_canvas_soft) # [0.0-BG, 1.0-stroke]
#### sampling part - start ####
self.curr_canvas_hard = curr_canvas_hard
rasterizor_st = NeuralRasterizorStep(
raster_size=self.hps.raster_size,
position_format=self.hps.position_format)
if self.hps.cropping_type == 'v3':
cropping_func = self.image_cropping_v3
# elif self.hps.cropping_type == 'v2':
# cropping_func = self.image_cropping
else:
raise Exception('Unknown cropping_type', self.hps.cropping_type)
if self.hps.pasting_type == 'v3':
pasting_func = self.image_pasting_v3
# elif self.hps.pasting_type == 'v2':
# pasting_func = self.image_pasting
else:
raise Exception('Unknown pasting_type', self.hps.pasting_type)
for time_i in range(self.hps.max_seq_len):
cursor_position_non_grad = tf.stop_gradient(cursor_position_loop) # (N, 1, 2), in size [0.0, 1.0)
curr_window_size = tf.multiply(prev_scaling, tf.stop_gradient(prev_window_size)) # float, with grad
curr_window_size = tf.maximum(curr_window_size, tf.cast(self.hps.min_window_size, tf.float32))
curr_window_size = tf.minimum(curr_window_size, tf.cast(image_size, tf.float32))
## patch-level encoding
# Here, we make the gradients from canvas_z to curr_canvas_hard be None to avoid recurrent gradient propagation.
curr_canvas_hard_non_grad = tf.stop_gradient(self.curr_canvas_hard)
curr_canvas_hard_non_grad = tf.expand_dims(curr_canvas_hard_non_grad, axis=-1)
# input_photo: (N, image_size, image_size, 1/3), [0.0-stroke, 1.0-BG]
crop_inputs = tf.concat([1.0 - input_photo, curr_canvas_hard_non_grad], axis=-1) # (N, H_p, W_p, 1/3+1)
cropped_outputs = cropping_func(cursor_position_non_grad, crop_inputs, image_size, curr_window_size)
index_offset = self.hps.input_channel - 1
curr_patch_inputs = cropped_outputs[:, :, :, 0:1 + index_offset] # [0.0-BG, 1.0-stroke]
curr_patch_canvas_hard_non_grad = cropped_outputs[:, :, :, 1 + index_offset:2 + index_offset]
# (N, raster_size, raster_size, 1), [0.0-BG, 1.0-stroke]
curr_patch_inputs = 1.0 - curr_patch_inputs # [0.0-stroke, 1.0-BG]
curr_patch_inputs = self.normalize_image_m1to1(curr_patch_inputs)
# (N, raster_size, raster_size, 1/3), [-1.0-stroke, 1.0-BG]
# Normalizing image
curr_patch_canvas_hard_non_grad = 1.0 - curr_patch_canvas_hard_non_grad # [0.0-stroke, 1.0-BG]
curr_patch_canvas_hard_non_grad = self.normalize_image_m1to1(curr_patch_canvas_hard_non_grad) # [-1.0-stroke, 1.0-BG]
## image-level encoding
combined_z = self.build_combined_encoder(
curr_patch_canvas_hard_non_grad,
curr_patch_inputs,
1.0 - curr_canvas_hard_non_grad,
input_photo,
cursor_position_non_grad,
image_size,
curr_window_size) # (N, z_size)
combined_z = tf.expand_dims(combined_z, axis=1) # (N, 1, z_size)
curr_window_size_top_side_norm_non_grad = \
tf.stop_gradient(curr_window_size / tf.cast(image_size, tf.float32))
curr_window_size_bottom_side_norm_non_grad = \
tf.stop_gradient(curr_window_size / tf.cast(self.hps.min_window_size, tf.float32))
if not self.hps.concat_win_size:
combined_z = tf.concat([tf.stop_gradient(prev_width), combined_z], 2) # (N, 1, 2+z_size)
else:
combined_z = tf.concat([tf.stop_gradient(prev_width),
curr_window_size_top_side_norm_non_grad,
curr_window_size_bottom_side_norm_non_grad,
combined_z],
2) # (N, 1, 2+z_size)
if self.hps.concat_cursor:
prev_input_x = tf.concat([cursor_position_non_grad, combined_z], 2) # (N, 1, 2+2+z_size)
else:
prev_input_x = combined_z # (N, 1, 2+z_size)
h_output, next_state = self.build_seq_decoder(self.dec_cell, prev_input_x, prev_state)
# h_output: (N * 1, n_out), next_state: (N, dec_rnn_size * 3)
[o_other_params, o_pen_ras] = self.get_mixture_coef(h_output)
# o_other_params: (N * 1, 6)
# o_pen_ras: (N * 1, 2), after softmax
o_other_params = tf.reshape(o_other_params, [-1, 1, 6]) # (N, 1, 6)
o_pen_ras_raw = tf.reshape(o_pen_ras, [-1, 1, 2]) # (N, 1, 2)
other_params_list.append(o_other_params)
pen_ras_list.append(o_pen_ras_raw)
#### sampling part - end ####
if self.hps.model_mode == 'train' or self.hps.model_mode == 'eval' or self.hps.model_mode == 'eval_sample':
# use renderer here to convert the strokes to image
curr_other_params = tf.squeeze(o_other_params, axis=1) # (N, 6), (x1, y1)=[0.0, 1.0], (x2, y2)=[-1.0, 1.0]
x1y1, x2y2, width2, scaling = curr_other_params[:, 0:2], curr_other_params[:, 2:4],\
curr_other_params[:, 4:5], curr_other_params[:, 5:6]
x0y0 = tf.zeros_like(x2y2) # (N, 2), [-1.0, 1.0]
x0y0 = tf.div(tf.add(x0y0, 1.0), 2.0) # (N, 2), [0.0, 1.0]
x2y2 = tf.div(tf.add(x2y2, 1.0), 2.0) # (N, 2), [0.0, 1.0]
widths = tf.concat([tf.squeeze(prev_width, axis=1), width2], axis=1) # (N, 2)
curr_other_params = tf.concat([x0y0, x1y1, x2y2, widths], axis=-1) # (N, 8), (x0, y0)&(x2, y2)=[0.0, 1.0]
curr_stroke_image = rasterizor_st.raster_func_stroke_abs(curr_other_params)
# (N, raster_size, raster_size), [0.0-BG, 1.0-stroke]
curr_stroke_image_large = pasting_func(cursor_position_loop, curr_stroke_image,
image_size, curr_window_size,
is_differentiable=self.hps.pasting_diff)
# (N, image_size, image_size), [0.0-BG, 1.0-stroke]
## soft
if not self.hps.use_softargmax:
curr_state_soft = o_pen_ras[:, 1:2] # (N, 1)
else:
curr_state_soft = self.differentiable_argmax(o_pen_ras, self.hps.soft_beta) # (N, 1)
curr_state_soft = tf.expand_dims(curr_state_soft, axis=1) # (N, 1, 1)
filter_curr_stroke_image_soft = tf.multiply(tf.subtract(1.0, curr_state_soft), curr_stroke_image_large)
# (N, image_size, image_size), [0.0-BG, 1.0-stroke]
curr_canvas_soft = tf.add(curr_canvas_soft, filter_curr_stroke_image_soft) # [0.0-BG, 1.0-stroke]
## hard
curr_state_hard = tf.expand_dims(tf.cast(tf.argmax(o_pen_ras_raw, axis=-1), dtype=tf.float32),
axis=-1) # (N, 1, 1)
filter_curr_stroke_image_hard = tf.multiply(tf.subtract(1.0, curr_state_hard), curr_stroke_image_large)
# (N, image_size, image_size), [0.0-BG, 1.0-stroke]
self.curr_canvas_hard = tf.add(self.curr_canvas_hard, filter_curr_stroke_image_hard) # [0.0-BG, 1.0-stroke]
self.curr_canvas_hard = tf.clip_by_value(self.curr_canvas_hard, 0.0, 1.0) # [0.0-BG, 1.0-stroke]
next_width = o_other_params[:, :, 4:5]
next_scaling = o_other_params[:, :, 5:6]
next_window_size = tf.multiply(next_scaling, tf.stop_gradient(curr_window_size)) # float, with grad
window_size_before_max_min = next_window_size # (N, 1, 1), large-level
win_size_before_max_min_list.append(window_size_before_max_min)
next_window_size = tf.maximum(next_window_size, tf.cast(self.hps.min_window_size, tf.float32))
next_window_size = tf.minimum(next_window_size, tf.cast(image_size, tf.float32))
prev_state = next_state
prev_width = next_width * curr_window_size / next_window_size # (N, 1, 1)
prev_scaling = next_scaling # (N, 1, 1))
prev_window_size = curr_window_size
# update the cursor position
new_cursor_offsets = tf.multiply(o_other_params[:, :, 2:4],
tf.divide(curr_window_size, 2.0)) # (N, 1, 2), window-level
new_cursor_offset_next = new_cursor_offsets
new_cursor_offset_next = tf.concat([new_cursor_offset_next[:, :, 1:2], new_cursor_offset_next[:, :, 0:1]], axis=-1)
cursor_position_loop_large = tf.multiply(cursor_position_loop, tf.cast(image_size, tf.float32))
if self.hps.stop_accu_grad:
stroke_position_next = tf.stop_gradient(cursor_position_loop_large) + new_cursor_offset_next # (N, 1, 2), large-level
else:
stroke_position_next = cursor_position_loop_large + new_cursor_offset_next # (N, 1, 2), large-level
stroke_position_before_max_min = stroke_position_next # (N, 1, 2), large-level
pos_before_max_min_list.append(stroke_position_before_max_min)
if self.hps.cursor_type == 'next':
cursor_position_loop_large = stroke_position_next # (N, 1, 2), large-level
else:
raise Exception('Unknown cursor_type')
cursor_position_loop_large = tf.maximum(cursor_position_loop_large, 0.0)
cursor_position_loop_large = tf.minimum(cursor_position_loop_large, tf.cast(image_size - 1, tf.float32))
cursor_position_loop = tf.div(cursor_position_loop_large, tf.cast(image_size, tf.float32))
curr_canvas_soft = tf.clip_by_value(curr_canvas_soft, 0.0, 1.0) # (N, raster_size, raster_size), [0.0-BG, 1.0-stroke]
other_params_ = tf.reshape(tf.concat(other_params_list, axis=1), [-1, 6]) # (N * max_seq_len, 6)
pen_ras_ = tf.reshape(tf.concat(pen_ras_list, axis=1), [-1, 2]) # (N * max_seq_len, 2)
pos_before_max_min_ = tf.concat(pos_before_max_min_list, axis=1) # (N, max_seq_len, 2)
win_size_before_max_min_ = tf.concat(win_size_before_max_min_list, axis=1) # (N, max_seq_len, 1)
return other_params_, pen_ras_, prev_state, curr_canvas_soft, curr_canvas_soft_rgb, \
pos_before_max_min_, win_size_before_max_min_
def differentiable_argmax(self, input_pen, soft_beta):
"""
Differentiable argmax trick.
:param input_pen: (N, n_class)
:return: pen_state: (N, 1)
"""
def sign_onehot(x):
"""
:param x: (N, n_class)
:return: (N, n_class)
"""
y = tf.sign(tf.reduce_max(x, axis=-1, keepdims=True) - x)
y = (y - 1) * (-1)
return y
def softargmax(x, beta=1e2):
"""
:param x: (N, n_class)
:param beta: 1e10 is the best. 1e2 is acceptable.
:return: (N)
"""
x_range = tf.cumsum(tf.ones_like(x), axis=1) # (N, 2)
return tf.reduce_sum(tf.nn.softmax(x * beta) * x_range, axis=1) - 1
## Better to use softargmax(beta=1e2). The sign_onehot's gradient is close to zero.
# pen_onehot = sign_onehot(input_pen) # one-hot form, (N * max_seq_len, 2)
# pen_state = pen_onehot[:, 1:2] # (N * max_seq_len, 1)
pen_state = softargmax(input_pen, soft_beta)
pen_state = tf.expand_dims(pen_state, axis=1) # (N * max_seq_len, 1)
return pen_state
def build_losses(self, target_sketch, pred_raster_imgs, pred_params,
pos_before_max_min, win_size_before_max_min, image_size):
def get_raster_loss(pred_imgs, gt_imgs, loss_type):
perc_layer_losses_raw = []
perc_layer_losses_weighted = []
perc_layer_losses_norm = []
if loss_type == 'l1':
ras_cost = tf.reduce_mean(tf.abs(tf.subtract(gt_imgs, pred_imgs))) # ()
elif loss_type == 'l1_small':
gt_imgs_small = tf.image.resize_images(tf.expand_dims(gt_imgs, axis=3), (32, 32))
pred_imgs_small = tf.image.resize_images(tf.expand_dims(pred_imgs, axis=3), (32, 32))
ras_cost = tf.reduce_mean(tf.abs(tf.subtract(gt_imgs_small, pred_imgs_small))) # ()
elif loss_type == 'mse':
ras_cost = tf.reduce_mean(tf.pow(tf.subtract(gt_imgs, pred_imgs), 2)) # ()
elif loss_type == 'perceptual':
return_map_pred = vgg_net_slim(pred_imgs, image_size)
return_map_gt = vgg_net_slim(gt_imgs, image_size)
perc_loss_type = 'l1' # [l1, mse]
weighted_map = {'ReLU1_1': 100.0, 'ReLU1_2': 100.0,
'ReLU2_1': 100.0, 'ReLU2_2': 100.0,
'ReLU3_1': 10.0, 'ReLU3_2': 10.0, 'ReLU3_3': 10.0,
'ReLU4_1': 1.0, 'ReLU4_2': 1.0, 'ReLU4_3': 1.0,
'ReLU5_1': 1.0, 'ReLU5_2': 1.0, 'ReLU5_3': 1.0}
for perc_layer in self.hps.perc_loss_layers:
if perc_loss_type == 'l1':
perc_layer_loss = tf.reduce_mean(tf.abs(tf.subtract(return_map_pred[perc_layer],
return_map_gt[perc_layer]))) # ()
elif perc_loss_type == 'mse':
perc_layer_loss = tf.reduce_mean(tf.pow(tf.subtract(return_map_pred[perc_layer],
return_map_gt[perc_layer]), 2)) # ()
else:
raise NameError('Unknown perceptual loss type:', perc_loss_type)
perc_layer_losses_raw.append(perc_layer_loss)
assert perc_layer in weighted_map
perc_layer_losses_weighted.append(perc_layer_loss * weighted_map[perc_layer])
for loop_i in range(len(self.hps.perc_loss_layers)):
perc_relu_loss_raw = perc_layer_losses_raw[loop_i] # ()
if self.hps.model_mode == 'train':
curr_relu_mean = (self.perc_loss_mean_list[loop_i] * self.last_step_num + perc_relu_loss_raw) / (self.last_step_num + 1.0)
relu_cost_norm = perc_relu_loss_raw / curr_relu_mean
else:
relu_cost_norm = perc_relu_loss_raw
perc_layer_losses_norm.append(relu_cost_norm)
perc_layer_losses_raw = tf.stack(perc_layer_losses_raw, axis=0)
perc_layer_losses_norm = tf.stack(perc_layer_losses_norm, axis=0)
if self.hps.perc_loss_fuse_type == 'max':
ras_cost = tf.reduce_max(perc_layer_losses_norm)
elif self.hps.perc_loss_fuse_type == 'add':
ras_cost = tf.reduce_mean(perc_layer_losses_norm)
elif self.hps.perc_loss_fuse_type == 'raw_add':
ras_cost = tf.reduce_mean(perc_layer_losses_raw)
elif self.hps.perc_loss_fuse_type == 'weighted_sum':
ras_cost = tf.reduce_mean(perc_layer_losses_weighted)
else:
raise NameError('Unknown perc_loss_fuse_type:', self.hps.perc_loss_fuse_type)
elif loss_type == 'triplet':
raise Exception('Solution for triplet loss is coming soon.')
else:
raise NameError('Unknown loss type:', loss_type)
if loss_type != 'perceptual':
for perc_layer_i in self.hps.perc_loss_layers:
perc_layer_losses_raw.append(tf.constant(0.0))
perc_layer_losses_norm.append(tf.constant(0.0))
perc_layer_losses_raw = tf.stack(perc_layer_losses_raw, axis=0)
perc_layer_losses_norm = tf.stack(perc_layer_losses_norm, axis=0)
return ras_cost, perc_layer_losses_raw, perc_layer_losses_norm
gt_raster_images = tf.squeeze(target_sketch, axis=3) # (N, raster_h, raster_w), [0.0-stroke, 1.0-BG]
raster_cost, perc_relu_losses_raw, perc_relu_losses_norm = \
get_raster_loss(pred_raster_imgs, gt_raster_images, loss_type=self.hps.raster_loss_base_type)
def get_stroke_num_loss(input_strokes):
ending_state = input_strokes[:, :, 0] # (N, seq_len)
stroke_num_loss_pre = tf.reduce_mean(ending_state) # larger is better, [0.0, 1.0]
stroke_num_loss = 1.0 - stroke_num_loss_pre # lower is better, [0.0, 1.0]
return stroke_num_loss
stroke_num_cost = get_stroke_num_loss(pred_params) # lower is better
def get_pos_outside_loss(pos_before_max_min_):
pos_after_max_min = tf.maximum(pos_before_max_min_, 0.0)
pos_after_max_min = tf.minimum(pos_after_max_min, tf.cast(image_size - 1, tf.float32)) # (N, max_seq_len, 2)
pos_outside_loss = tf.reduce_mean(tf.abs(pos_before_max_min_ - pos_after_max_min))
return pos_outside_loss
pos_outside_cost = get_pos_outside_loss(pos_before_max_min) # lower is better
def get_win_size_outside_loss(win_size_before_max_min_, min_window_size):
win_size_outside_top_loss = tf.divide(
tf.maximum(win_size_before_max_min_ - tf.cast(image_size, tf.float32), 0.0),
tf.cast(image_size, tf.float32)) # (N, max_seq_len, 1)
win_size_outside_bottom_loss = tf.divide(
tf.maximum(tf.cast(min_window_size, tf.float32) - win_size_before_max_min_, 0.0),
tf.cast(min_window_size, tf.float32)) # (N, max_seq_len, 1)
win_size_outside_loss = tf.reduce_mean(win_size_outside_top_loss + win_size_outside_bottom_loss)
return win_size_outside_loss
win_size_outside_cost = get_win_size_outside_loss(win_size_before_max_min, self.hps.min_window_size) # lower is better
def get_early_pen_states_loss(input_strokes, curr_start, curr_end):
# input_strokes: (N, max_seq_len, 7)
pred_early_pen_states = input_strokes[:, curr_start:curr_end, 0] # (N, curr_early_len)
pred_early_pen_states_min = tf.reduce_min(pred_early_pen_states, axis=1) # (N), should not be 1
early_pen_states_loss = tf.reduce_mean(pred_early_pen_states_min) # lower is better
return early_pen_states_loss
early_pen_states_cost = get_early_pen_states_loss(pred_params,
self.early_pen_loss_start_idx, self.early_pen_loss_end_idx)
return raster_cost, stroke_num_cost, pos_outside_cost, win_size_outside_cost, \
early_pen_states_cost, \
perc_relu_losses_raw, perc_relu_losses_norm
def build_training_op_split(self, raster_cost, sn_cost, cursor_outside_cost, win_size_outside_cost,
early_pen_states_cost):
total_cost = self.hps.raster_loss_weight * raster_cost + \
self.hps.early_pen_loss_weight * early_pen_states_cost + \
self.stroke_num_loss_weight * sn_cost + \
self.hps.outside_loss_weight * cursor_outside_cost + \
self.hps.win_size_outside_loss_weight * win_size_outside_cost
tvars = [var for var in tf.trainable_variables()
if 'raster_unit' not in var.op.name and 'VGG16' not in var.op.name]
gvs = self.optimizer.compute_gradients(total_cost, var_list=tvars)
return total_cost, gvs
def build_training_op(self, grad_list):
with tf.variable_scope('train_op', reuse=tf.AUTO_REUSE):
gvs = self.average_gradients(grad_list)
g = self.hps.grad_clip
for grad, var in gvs:
print('>>', var.op.name)
if grad is None:
print(' >> None value')
capped_gvs = [(tf.clip_by_value(grad, -g, g), var) for grad, var in gvs]
self.train_op = self.optimizer.apply_gradients(
capped_gvs, global_step=self.global_step, name='train_step')
def average_gradients(self, grads_list):
"""
Compute the average gradients.
:param grads_list: list(of length N_GPU) of list(grad, var)
:return:
"""
avg_grads = []
for grad_and_vars in zip(*grads_list):
grads = []
for g, _ in grad_and_vars:
expanded_g = tf.expand_dims(g, 0)
grads.append(expanded_g)
grad = tf.concat(grads, axis=0)
grad = tf.reduce_mean(grad, axis=0)
v = grad_and_vars[0][1]
grad_and_var = (grad, v)
avg_grads.append(grad_and_var)
return avg_grads