def show_sample(self,
content_img,
style_img,
concate=True,
denorm=True,
deprocess=True):
gen_img = self.generate(content_img, style_img)
if concate:
return utils.show_images(
np.concatenate([content_img, style_img, gen_img]), denorm,
deprocess)
if denorm:
content_img = utils.de_norm(content_img)
style_img = utils.de_norm(style_img)
gen_img = utils.de_norm(gen_img)
if deprocess:
content_img = utils.deprocess(content_img)
style_img = utils.deprocess(style_img)
gen_img = utils.deprocess(gen_img)
cv2_imshow(content_img[0])
cv2_imshow(style_img[0])
cv2_imshow(gen_img[0])
def deep_dream(image, model, iterations, lr, octave_scale, num_octaves):
""" Main deep dream method """
image = preprocess(image).unsqueeze(0).to(device).data.numpy()
# Extract image representations for each octave
octaves = [image]
for _ in range(num_octaves - 1):
octaves.append(
nd.zoom(octaves[-1], (1, 1, 1 / octave_scale, 1 / octave_scale),
order=1))
detail = np.zeros_like(octaves[-1])
for octave, octave_base in enumerate(
tqdm.tqdm(octaves[::-1], desc="Dreaming")):
if octave > 0:
# Upsample detail to new octave dimension
detail = nd.zoom(detail,
np.array(octave_base.shape) /
np.array(detail.shape),
order=1)
# Add deep dream detail from previous octave to new base
input_image = octave_base + detail
# Get new deep dream image
dreamed_image = dream(input_image, model, iterations, lr)
# Extract deep dream details
detail = dreamed_image - octave_base
return deprocess(dreamed_image)
def deep_dream(image, model, iterations, lr, octave_scale, num_octaves):
""" Main deep dream method """
image = preprocess(image).unsqueeze(0).cpu().data.numpy()
# Extract octaves for each dimension
octaves = [image]
for i in range(num_octaves - 1):
octaves.append(
nd.zoom(octaves[-1],
(1, 1, 1.0 / octave_scale, 1.0 / octave_scale),
order=1))
detail = np.zeros_like(octaves[-1])
for octave, octave_base in enumerate(
tqdm.tqdm(octaves[::-1], desc="Dreaming")):
if octave > 0:
# Upsample detail to new dimension
h, w = octave_base.shape[-2:]
dh, dw = detail.shape[-2:]
detail = nd.zoom(detail, (1, 1, h / dh, w / dw), order=1)
# Add deep dream detail from previous octave to new base
input_image = octave_base + detail
# Get new deep dream image
dreamed_image = dream(input_image, model, iterations, lr)
# Extract deep dream details
detail = dreamed_image - octave_base
return deprocess(dreamed_image)
def deep_dream(image, model, iterations, lr, octave_scale, num_octaves):
""" Main deep dream method """
image = preprocess(image).unsqueeze(0).cuda()
# Extract image representations for each octave
octaves = [image]
for _ in range(num_octaves - 1):
image = nn.functional.interpolate(image,
scale_factor=1 / octave_scale,
mode='bilinear',
align_corners=False)
octaves.append(image)
detail = torch.zeros(octaves[-1].size(), dtype=torch.float).cuda()
for octave, octave_base in enumerate(
tqdm.tqdm(octaves[::-1], desc="Dreaming")):
if octave > 0:
# Upsample detail to new octave dimension
detail = nn.functional.interpolate(detail,
size=octave_base.size()[-2:],
mode='bilinear',
align_corners=False)
# Add deep dream detail from previous octave to new base
input_image = octave_base + detail
# Get new deep dream image
dreamed_image = dream(input_image.cuda(), model, iterations, lr)
# Extract deep dream details
detail = dreamed_image - octave_base
return deprocess(dreamed_image.cpu().data.numpy())
def main():
# Load model
model = load_model(MODEL_PATH, custom_objects={'AdaIN': AdaIN})
# Get content image
content = get_image(CONTENT_PATH, resize=False)
content = preprocess(content)
content = np.expand_dims(content, axis=0)
# Get style image
style = get_image(STYLE_PATH, resize=False)
style = preprocess(style)
style = np.expand_dims(style, axis=0)
# Set alpha Value
alpha = tf.convert_to_tensor(ALPHA) # 0 < alpha <= 1
alpha = np.expand_dims(alpha, axis=0)
# Do inference
y = model.predict([content, style, alpha])[0]
# Convert output array to image
y = np.squeeze(y, axis=0)
y = deprocess(y)
img = array_to_img(y)
# Show image
img.show(command='fim')
def show_imgs(self, img):
if len(img.shape) == 4:
return utils.show_images(img, self.normalize, self.preprocessing)
if self.normalize:
img = utils.de_norm(img)
if self.preprocessing:
img = utils.deprocess(img)
cv2_imshow(img)
def makeup(self):
if self.no_face:
return self.face
else:
self.face = cv2.cvtColor(self.face, cv2.COLOR_RGB2BGR)
no_makeup = cv2.resize(self.face, (self.img_size, self.img_size))
X_img = np.expand_dims(preprocess(no_makeup), 0)
Xs_ = self.sess.run(self.Xs, feed_dict={self.X: X_img, self.Y: self.Y_img})
Xs_ = deprocess(Xs_)
return cv2.cvtColor(cv2.resize(Xs_[0]*255, (return_img_size, return_img_size)), cv2.COLOR_RGB2BGR)
def optimize_single(seed_fname, guide_fname, model_name, layer, iter_n,
max_thres, net, verbose=True):
seed = utils.read_image_rgb(seed_fname, net)
guide = utils.read_image_rgb(guide_fname, net)
objective = get_layer_objective(layer, net)
ad, bc, cd = constOptimize(
net, seed, guide, iter_n=iter_n, max_thres=max_thres,
end=layer, objective=objective, verbose=verbose)
print cd
return utils.deprocess(net, ad), bc, cd
def save_img(self, img, suffix, iter_):
"""Summary
Args:
img (numpy array): Output Image
suffix (string): filename suffix
iter_ (integer): the iteration value
"""
img = deprocess(img)
img = np.clip(img, 0, 1)
file_name = self.img_list[self.config["start_position"] + iter_]
file_name = file_name.split("/")[-1]
plt.imsave(self.outpath + "/{}{}".format(suffix, file_name), img)
def loss(y_true, y_pred):
# y_true == input == [content, style]
out, adain = y_pred[0], y_pred[1]
# Encode output and compute content_loss
out = deprocess(out)
out = preprocess(out)
enc_out = encoder(out)
content_loss = tf.reduce_sum(tf.reduce_mean(tf.square(enc_out - adain), axis=[1, 2]))
# Compute style loss from vgg relus
style = y_true[1]
style_featues = vgg19_relus(style)
gen_features = vgg19_relus(out)
style_layer_loss = []
for enc_style_feat, enc_gen_feat in zip(style_featues, gen_features):
meanS, varS = tf.nn.moments(enc_style_feat, [1, 2])
meanG, varG = tf.nn.moments(enc_gen_feat, [1, 2])
sigmaS = tf.sqrt(varS + epsilon)
sigmaG = tf.sqrt(varG + epsilon)
l2_mean = tf.reduce_sum(tf.square(meanG - meanS))
l2_sigma = tf.reduce_sum(tf.square(sigmaG - sigmaS))
style_layer_loss.append(l2_mean + l2_sigma)
style_loss = tf.reduce_sum(style_layer_loss)
# Compute color loss
style_color_mean, style_color_var = tf.nn.moments(style, [1, 2])
gen_color_mean, gen_color_var = tf.nn.moments(out, [1, 2])
color_sigmaS = tf.sqrt(style_color_var)
color_sigmaG = tf.sqrt(gen_color_var)
l2_mean = tf.reduce_sum(tf.square(gen_color_mean - style_color_mean))
l2_sigma = tf.reduce_sum(tf.square(color_sigmaG - color_sigmaS))
color_loss = l2_mean + l2_sigma
# Compute the total loss
weighted_style_loss = style_weight * style_loss
weighted_color_loss = color_loss * color_weight
total_loss = content_loss + weighted_style_loss + weighted_color_loss
return total_loss, content_loss, weighted_style_loss, weighted_color_loss
def optimize_single(seed_fname,
guide_fname,
model_name,
layer,
iter_n,
max_thres,
net,
verbose=True):
seed = utils.read_image_rgb(seed_fname, net)
guide = utils.read_image_rgb(guide_fname, net)
objective = get_layer_objective(layer, net)
ad, bc, cd = constOptimize(net,
seed,
guide,
iter_n=iter_n,
max_thres=max_thres,
end=layer,
objective=objective,
verbose=verbose)
print cd
return utils.deprocess(net, ad), bc, cd
def train(self, iterations, frequency, target_path):
with self.styling_sess.as_default(), self.styling_g.as_default():
# Initialize variables in vgg16 model (dropout etc.)
self.styling_sess.run(tf.initialize_all_variables())
for it_i in range(iterations):
# Run optimizer without dropout to make model more deterministic
utils.run_without_dropout(self.styling_sess, self.optimizer)
# Print step of gradient descent
print("%d" % it_i)
# Save generated image to target path with required frequency
if it_i % frequency == 0:
generated_image4D = utils.run_without_dropout(self.styling_sess, self.styled_image)
# Choose the first image in the batch of 3D images: generated_image4D
# (there is only one image in the batch), deprocess it, and save the image
Image.fromarray(utils.deprocess(generated_image4D[0])).save(target_path + "img-%d.png" % it_i)
def deep_dream(image, model, output=None, iterations=20, lr=0.01, octave_scale=1.4, num_octaves=10, skip_octaves=0):
"""Main deep dream method"""
image = preprocess(image).unsqueeze(0).cpu().data.numpy()
# Generate scaled images for ech octave
octaves = [image]
for i in range(num_octaves - 1):
if i >= skip_octaves:
octaves.append(nd.zoom(octaves[-1], (1, 1, 1/octave_scale, 1/octave_scale), order=1))
detail = np.zeros_like(octaves[-1])
for octave, octave_base in enumerate(octaves[::-1]):
if octave > 0:
# Upsample to match octave dimension
detail = nd.zoom(detail, np.array(octave_base.shape)/np.array(detail.shape), order=1)
# Add detail to base image
input_image = octave_base + detail
# Run dream process
dreamed_image = dream(input_image, model, output, iterations, lr)
# Remove base image to obtain purely dreamed details
detail = dreamed_image - octave_base
return deprocess(dreamed_image)
guide_img_tensor, config.small_img_num, config.small_zoom_ratio)
# Training
dreaming_detail_np = np.zeros_like(input_small_imgs_np_list[-1])
output_img_np = dreaming_detail_np
for index, base_img_np in enumerate(
tqdm.tqdm(input_small_imgs_np_list[::-1], desc="Processing")):
if index > 0:
zoom = np.array(base_img_np.shape) / np.array(
dreaming_detail_np.shape)
dreaming_detail_np = utils.zooming_img(dreaming_detail_np, zoom)
input_img_np = base_img_np + dreaming_detail_np
iteration = (config.each_iteration +
10) if (index == len(input_small_imgs_np_list) -
1) else config.each_iteration
guide_img_np = guide_small_imgs_np_list[len(input_small_imgs_np_list) -
1 - index]
output_img_np = dream_process(input_img_np, guide_img_np,
deep_dreamer_net, config.lr, config.mode,
iteration, index, device)
dreaming_detail_np = output_img_np - base_img_np
# Display
dreamed_img = utils.deprocess(output_img_np)
plt.figure()
plt.imshow(dreamed_img)
plt.show()
plt.imsave(config.output_img_path, dreamed_img)
def main():
"""Create the model and start the training."""
'''
1. create image reader
'''
with tf.device('/cpu:0'):
left, right, count, batch_size = load_data([a.left_dir, a.right_dir], [a.height, a.width], name='load_data')
if a.is_val:
left_val, right_val, val_count, val_batch_size = load_data([a.left_val_dir, a.right_val_dir], [a.height, a.width], name='load_val_data')
print('Num_data: {}'.format(count))
if a.is_val:
print('Num_val: {}'.format(val_count))
'''
2. build model, the prediction
'''
with tf.device('/gpu:0'):
with tf.name_scope('build_graph'):
loss, l_init, l_ref, left_initial_disp, right_initial_disp, left_refined_disp, right_refined_disp = build_model([left, right], is_train=True, reuse=False)
if a.is_val:
val_loss, _, _, _, _, _, _ = build_model([left_val, right_val], is_train=False, reuse=True)
'''
3. do updating
'''
with tf.name_scope('train'):
global_step = tf.Variable(0, trainable=False, name='global_step')
# lr = tf.train.exponential_decay(LEARNING_RATE, global_step, 10, 0.96, staircase=True, name='learning_rate')
rate = tf.pow(0.5, tf.cast(tf.cast(global_step/a.schedule_freq, tf.int32), tf.float32))
lr = a.lr * rate
tf.summary.scalar('learning_rate', lr, collections=['train_summary'])
tf.summary.scalar('step', global_step, collections=['train_summary'])
update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
with tf.control_dependencies(update_ops):
optimizer = tf.train.AdamOptimizer(lr)
optimize = optimizer.minimize(loss, global_step)
with tf.name_scope('loss'):
tf.summary.scalar('loss', loss, collections=['train_summary'])
tf.summary.scalar('loss_init', l_init, collections=['train_summary'])
tf.summary.scalar('loss_ref', l_ref, collections=['train_summary'])
loss_val = tf.placeholder(tf.float32, [])
loss_val_sum = tf.summary.scalar('loss_val', loss_val)
with tf.name_scope('input_images'):
tf.summary.image("left", deprocess(left), max_outputs=1, collections=['train_summary'])
tf.summary.image("right", deprocess(right), max_outputs=1, collections=['train_summary'])
with tf.name_scope('disp_images'):
tf.summary.image("left_disp_refined", left_refined_disp, max_outputs=1, collections=['train_summary'])
tf.summary.image("right_disp_refined", right_refined_disp, max_outputs=1, collections=['train_summary'])
tf.summary.image("left_disp_init", left_initial_disp, max_outputs=1, collections=['train_summary'])
tf.summary.image("right_disp_init", right_initial_disp, max_outputs=1, collections=['train_summary'])
'''
3. training setting
'''
with tf.name_scope('save'):
saver = tf.train.Saver(max_to_keep=8)
summary_writer = tf.summary.FileWriter(a.summary_dir)
# summary_op = tf.summary.merge([loss_sum,left_sum,right_sum,left_disp_sum,right_disp_sum, step_sum, lr_sum])
summary_op = tf.summary.merge_all(key='train_summary')
init = tf.global_variables_initializer()
'''
4. begin to train
'''
config = tf.ConfigProto(allow_soft_placement=True)
with tf.Session(config=config) as sess:
if a.resume_dir is not None:
# restoring from the checkpoint file
ckpt = tf.train.get_checkpoint_state(a.resume_dir)
if ckpt is not None:
tf.train.Saver().restore(sess, ckpt.model_checkpoint_path)
sess.run(global_step.assign(0)) # reset global_step to zero
print('Reload from: {}'.format(a.resume_dir))
else:
sess.run(init)
else:
sess.run(init)
# sess.run(load_pretrained_parameters)
summary_writer.add_graph(sess.graph)
tf.train.start_queue_runners(sess=sess)
for step in range(a.num_steps):
_, l, step = sess.run([optimize, loss, global_step])
train_epoch = step * a.batch_size // count
if step % a.summary_freq == 0:
s = sess.run(summary_op)
summary_writer.add_summary(s, step)
summary_writer.flush()
print('-------- summary saved --------')
if a.is_val:
if step % count == 0:
print('Running Validation')
# iterate through validation set
total_vl = 0
for i in range(0, val_count):
vl = sess.run(val_loss)
total_vl = total_vl + vl
vl_avg = 1.0*total_vl/val_count
s = sess.run(loss_val_sum, {loss_val: vl_avg})
summary_writer.add_summary(s, step)
summary_writer.flush()
print('-------- training_loss:{0:.4f} validation_loss:{1:.4f}'.format(l, vl_avg))
if step % a.save_freq == 0 and step != 0:
saver.save(sess, a.checkpoint_dir + 'model.ckpt', global_step=step)
print('-------- checkpoint saved:{} --------'.format(step))
if step % a.print_summary_freq == 0:
print('epoch:{0} step:{1} loss:{2:.4f}'.format(train_epoch, step, l))
# after loop
saver.save(sess, a.checkpoint_dir + 'model.ckpt', global_step=step)
print('-------- checkpoint saved:{} --------'.format(step))
def main():
if tf.__version__ != "1.0.0":
raise Exception("Tensorflow version 1.0.0 required")
if a.seed is None:
a.seed = random.randint(0, 2**31 - 1)
tf.set_random_seed(a.seed)
np.random.seed(a.seed)
random.seed(a.seed)
if not os.path.exists(a.output_dir):
os.makedirs(a.output_dir)
if a.mode == "test" or a.mode == "export":
if a.checkpoint is None:
raise Exception("checkpoint required for test mode")
# load some options from the checkpoint
options = {"which_direction", "ngf", "ndf", "lab_colorization"}
with open(os.path.join(a.checkpoint, "options.json")) as f:
for key, val in json.loads(f.read()).items():
if key in options:
print("loaded", key, "=", val)
setattr(a, key, val)
# disable these features in test mode
a.scale_size = CROP_SIZE
a.flip = False
for k, v in a._get_kwargs():
print(k, "=", v)
with open(os.path.join(a.output_dir, "options.json"), "w") as f:
f.write(json.dumps(vars(a), sort_keys=True, indent=4))
if a.mode == "export":
# export the generator to a meta graph that can be imported later for standalone generation
if a.lab_colorization:
raise Exception("export not supported for lab_colorization")
input = tf.placeholder(tf.string, shape=[1])
input_data = tf.decode_base64(input[0])
input_image = tf.image.decode_png(input_data)
# remove alpha channel if present
input_image = input_image[:, :, :3]
input_image = tf.image.convert_image_dtype(input_image,
dtype=tf.float32)
input_image.set_shape([CROP_SIZE, CROP_SIZE, 3])
batch_input = tf.expand_dims(input_image, axis=0)
with tf.variable_scope("generator") as scope:
batch_output = deprocess(
create_generator_Unet(a, preprocess(batch_input), 3))
output_image = tf.image.convert_image_dtype(batch_output,
dtype=tf.uint8)[0]
output_data = tf.image.encode_png(output_image)
output = tf.convert_to_tensor([tf.encode_base64(output_data)])
key = tf.placeholder(tf.string, shape=[1])
inputs = {"key": key.name, "input": input.name}
tf.add_to_collection("inputs", json.dumps(inputs))
outputs = {
"key": tf.identity(key).name,
"output": output.name,
}
tf.add_to_collection("outputs", json.dumps(outputs))
init_op = tf.global_variables_initializer()
restore_saver = tf.train.Saver()
export_saver = tf.train.Saver()
# config = tf.ConfigProto()
# config.allow_soft_placement = False
# config.log_device_placement = True
# config.gpu_options.allow_growth = True
# config.gpu_options.per_process_gpu_memory_fraction = 0.4
with tf.Session() as sess:
sess.run(init_op)
print("loading model from checkpoint")
checkpoint = tf.train.latest_checkpoint(a.checkpoint)
restore_saver.restore(sess, checkpoint)
print("exporting model")
export_saver.export_meta_graph(
filename=os.path.join(a.output_dir, "export.meta"))
export_saver.save(sess,
os.path.join(a.output_dir, "export"),
write_meta_graph=False)
return
examples = load_examples(a)
print("examples count = %d" % examples.count)
# inputs and targets are [batch_size, height, width, channels]
model = create_model(examples.inputs, examples.targets)
# undo colorization splitting on images that we use for display/output
if a.lab_colorization:
if a.which_direction == "AtoB":
# inputs is brightness, this will be handled fine as a grayscale image
# need to augment targets and outputs with brightness
targets = augment(examples.targets, examples.inputs)
outputs = augment(model.outputs, examples.inputs)
# inputs can be deprocessed normally and handled as if they are single channel
# grayscale images
inputs = deprocess(examples.inputs)
elif a.which_direction == "BtoA":
# inputs will be color channels only, get brightness from targets
inputs = augment(examples.inputs, examples.targets)
targets = deprocess(examples.targets)
outputs = deprocess(model.outputs)
else:
raise Exception("invalid direction")
else:
inputs = deprocess(examples.inputs)
targets = deprocess(examples.targets)
outputs = deprocess(model.outputs)
def convert(image):
if a.aspect_ratio != 1.0:
# upscale to correct aspect ratio
size = [CROP_SIZE, int(round(CROP_SIZE * a.aspect_ratio))]
image = tf.image.resize_images(
image, size=size, method=tf.image.ResizeMethod.BICUBIC)
return tf.image.convert_image_dtype(image,
dtype=tf.uint8,
saturate=True)
# reverse any processing on images so they can be written to disk or displayed to user
with tf.name_scope("convert_inputs"):
converted_inputs = convert(inputs)
with tf.name_scope("convert_targets"):
converted_targets = convert(targets)
with tf.name_scope("convert_outputs"):
converted_outputs = convert(outputs)
with tf.name_scope("encode_images"):
display_fetches = {
"paths":
examples.paths,
"inputs":
tf.map_fn(tf.image.encode_png,
converted_inputs,
dtype=tf.string,
name="input_pngs"),
"targets":
tf.map_fn(tf.image.encode_png,
converted_targets,
dtype=tf.string,
name="target_pngs"),
"outputs":
tf.map_fn(tf.image.encode_png,
converted_outputs,
dtype=tf.string,
name="output_pngs"),
}
# summaries
with tf.name_scope("inputs_summary"):
tf.summary.image("inputs", converted_inputs)
with tf.name_scope("targets_summary"):
tf.summary.image("targets", converted_targets)
with tf.name_scope("outputs_summary"):
tf.summary.image("outputs", converted_outputs)
with tf.name_scope("predict_real_summary"):
tf.summary.image(
"predict_real",
tf.image.convert_image_dtype(model.predict_real, dtype=tf.uint8))
with tf.name_scope("predict_fake_summary"):
tf.summary.image(
"predict_fake",
tf.image.convert_image_dtype(model.predict_fake, dtype=tf.uint8))
tf.summary.scalar("discriminator_loss", model.discrim_loss)
tf.summary.scalar("gennerator_loss", model.gen_loss)
tf.summary.scalar("generator_loss_GAN", model.gen_loss_GAN)
for var in tf.trainable_variables():
tf.summary.histogram(var.op.name + "/values", var)
for grad, var in model.discrim_grads_and_vars + model.gen_grads_and_vars:
tf.summary.histogram(var.op.name + "/gradients", grad)
with tf.name_scope("parameter_count"):
parameter_count = tf.reduce_sum(
[tf.reduce_prod(tf.shape(v)) for v in tf.trainable_variables()])
saver = tf.train.Saver(max_to_keep=1)
logdir = a.output_dir if (a.trace_freq > 0 or a.summary_freq > 0) else None
sv = tf.train.Supervisor(logdir=logdir, save_summaries_secs=0, saver=None)
with sv.managed_session() as sess:
print("parameter_count =", sess.run(parameter_count))
if a.checkpoint is not None:
print("loading model from checkpoint")
checkpoint = tf.train.latest_checkpoint(a.checkpoint)
saver.restore(sess, checkpoint)
max_steps = 2**32
if a.max_epochs is not None:
max_steps = examples.steps_per_epoch * a.max_epochs
if a.max_steps is not None:
max_steps = a.max_steps
if a.mode == "test":
# testing
# at most, process the test data once
test_start = time.time()
max_steps = min(examples.steps_per_epoch, max_steps)
for step in range(max_steps):
results = sess.run(display_fetches)
filesets = save_images(a, results)
for i, f in enumerate(filesets):
print("evaluated image", f["name"])
index_path = append_index(a, filesets)
with open(a.output_dir + '/test.time', 'w') as f:
f.write(str(time.time() - test_start) + '\n')
print("wrote index at", index_path)
else:
# training
start = time.time()
with open(a.output_dir + '/train.precession', 'a') as f:
f.write('step RFE ACC \n')
for step in range(max_steps):
def should(freq):
return freq > 0 and ((step + 1) % freq == 0
or step == max_steps - 1)
options = None
run_metadata = None
if should(a.trace_freq):
options = tf.RunOptions(
trace_level=tf.RunOptions.FULL_TRACE)
run_metadata = tf.RunMetadata()
fetches = {
"train": model.train,
"global_step": sv.global_step,
}
if should(a.progress_freq):
fetches["discrim_loss"] = model.discrim_loss
fetches["gen_loss"] = model.gen_loss
fetches["gen_loss_GAN"] = model.gen_loss_GAN
if should(a.summary_freq):
fetches["summary"] = sv.summary_op
if should(a.display_freq):
fetches["display"] = display_fetches
results = sess.run(fetches,
options=options,
run_metadata=run_metadata)
if should(a.summary_freq):
print("recording summary")
sv.summary_writer.add_summary(results["summary"],
results["global_step"])
if should(a.display_freq):
print("saving display images")
filesets = save_images(a,
results["display"],
step=results["global_step"])
append_index(a, filesets, step=True)
if should(a.trace_freq):
print("recording trace")
sv.summary_writer.add_run_metadata(
run_metadata, "step_%d" % results["global_step"])
if should(a.progress_freq):
# global_step will have the correct step count if we resume from a checkpoint
train_epoch = math.ceil(results["global_step"] /
examples.steps_per_epoch)
train_step = (results["global_step"] -
1) % examples.steps_per_epoch + 1
rate = (step + 1) * a.batch_size / (time.time() - start)
remaining = (max_steps - step) * a.batch_size / rate
print(
"progress epoch %d step %d image/sec %0.1f remaining %dm"
% (train_epoch, train_step, rate, remaining / 60))
print("discrim_loss", results["discrim_loss"])
print("gen_loss", results["gen_loss"])
print("gen_loss_GAN", results["gen_loss_GAN"])
#with open(a.output_dir + '/train.log', 'a') as f:
# f.write('step gen_loss discrim_loss')
if should(a.save_freq):
print("saving model")
saver.save(sess,
os.path.join(a.output_dir, "model"),
global_step=sv.global_step)
line_str = '%s %s %s\n' % (results["global_step"],
results["gen_loss"],
results["discrim_loss"])
with open(a.output_dir + '/train.log', 'a') as f:
f.write(line_str)
if should(a.test_freq):
if a.output_dir is None:
raise Exception("checkpoint required for test mode")
# load some options from the checkpoint
options = {
"which_direction", "ngf", "ndf", "lab_colorization"
}
with open(os.path.join(a.output_dir, "options.json")) as f:
for key, val in json.loads(f.read()).items():
if key in options:
print("loaded", key, "=", val)
setattr(a, key, val)
# disable these features in test mode
a.scale_size = CROP_SIZE
a.flip = False
# testing
# at most, process the test data once
max_steps = min(examples.steps_per_epoch, max_steps)
for step in range(max_steps):
test_results = sess.run(display_fetches)
test_filesets = save_images(
a, test_results, step=results["global_step"])
index_path = append_index(a, test_filesets)
if 'ThreeObj_gamma_1.0' in test_filesets[0]['outputs']:
#pdb.set_trace()
print(
'file',
'checkpoints.lr.%s/simulation_test/ossgan_sgan_l1/%s/images/%s'
%
(a.lr, a.f_type, test_filesets[0]['outputs']))
outp = cv2.imread(
'checkpoints.lr.%s/simulation_test/ossgan_sgan_l1/%s/images/%s'
%
(a.lr, a.f_type, test_filesets[0]['outputs']))
targ = cv2.imread(
'checkpoints.lr.%s/simulation_test/ossgan_sgan_l1/%s/images/%s'
%
(a.lr, a.f_type, test_filesets[0]['targets']))
inp = cv2.imread(
'checkpoints.lr.%s/simulation_test/ossgan_sgan_l1/%s/images/%s'
% (a.lr, a.f_type, test_filesets[0]['inputs']))
rfe = region_fitting_error(outp, targ)
acc = calculate_accuracy(inp, outp, targ)
#print('region fitting error', rfe)
#print('accuracy', acc)
line_str = '%s %s %s\n' % (results["global_step"],
rfe, acc)
with open(a.output_dir + '/train.precession',
'a') as f:
f.write(line_str)
print("wrote index at", index_path)
if sv.should_stop():
break
with open(a.output_dir + '/train.time', 'w') as f:
f.write(str(time.time() - start) + '\n')
with torch.no_grad():
for i in range(len(content_paths)):
content = load_image(content_paths[i])
content = preprocess(content, args.content_size)
content = content.to(device)
for j in range(len(style_paths)):
style = load_image(style_paths[j])
style = preprocess(style, args.style_size)
style = style.to(device)
if args.synthesis == 0:
output = style_transfer(content, style)
output = deprocess(output)
if len(content_paths) == 1 and len(style_paths) == 1:
# used a single content and style image
save_path = '%s/%s.%s' % (args.save_dir, args.save_name,
args.save_ext)
else:
# used a batch of content and style images
save_path = '%s/%s_%s.%s' % (args.save_dir, i, j,
args.save_ext)
print('Output image saved at:', save_path)
output.save(save_path)
else:
content = torch.rand(*content.shape).uniform_(0, 1).to(device)
for iteration in range(3):
def test_step(batch):
prediction = it_network(batch, training=False)
#prediction_norm = np.array(tf.clip_by_value(prediction, 0, 1)*255, dtype=np.uint8) # Poor quality, no convergence
#prediction_norm = np.array(tf.clip_by_value(prediction, 0, 255), dtype=np.uint8)
return deprocess(prediction)
def main():
purge()
global_step = tf.train.get_or_create_global_step()
input_node = tf.placeholder(
name='input_images',
shape=[1, None, None, 3],
dtype=tf.float32)
target_node = tf.placeholder(
name='target_images',
shape=[1, None, None, 3],
dtype=tf.float32)
is_train = tf.placeholder_with_default(False, (), name='is_training')
lr_image_paths = tf.placeholder(tf.string, shape=(None,), name='lr_image_paths')
hr_image_paths = tf.placeholder(tf.string, shape=(None,), name='hr_image_paths')
lr_valid_image = cv2.imread('dataset/lr_valid/jason.jpg')
lr_valid_image = cv2.cvtColor(lr_valid_image, cv2.COLOR_BGR2RGB)
lr_valid_image = np.expand_dims(lr_valid_image, axis=0)
lr_valid_image = lr_valid_image / 255
hr_valid_image = cv2.imread('dataset/hr_valid/jason.jpg')
hr_valid_image = cv2.cvtColor(hr_valid_image, cv2.COLOR_BGR2RGB)
hr_valid_image = (hr_valid_image / 255) * 2 - 1
hr_valid_image = np.expand_dims(hr_valid_image, axis=0)
with tf.variable_scope('load_image'):
input_image_lr = tf.map_fn(_parse_lr, lr_image_paths, dtype=tf.float32)
input_image_hr = tf.map_fn(_parse_hr, hr_image_paths, dtype=tf.float32)
# Normalize the low resolution image to [0, 1], high resolution to [-1, 1]
inputs_batch = preprocessLR(input_image_lr)
targets_batch = preprocess(input_image_hr)
with tf.name_scope('train'):
net_train = Generator(inputs_batch, is_train)
gen_output = net_train.arch_output
gen_output.set_shape([BATCH_SIZE, INPUT_SIZE[0] * 4, INPUT_SIZE[1] * 4, 3])
with tf.name_scope('valid'):
net_valid = Generator(input_node, tf.constant(False), reuse=True)
gen_valid = net_valid.arch_output
gen_valid.set_shape([1, INPUT_SIZE[0] * 4, INPUT_SIZE[1] * 4, 3])
gen_valid = deprocess(gen_valid)
valid_diff = gen_valid - target_node
valid_loss = tf.reduce_mean(tf.reduce_sum(tf.square(valid_diff), axis=[3]))
converted_outputs = tf.image.convert_image_dtype(gen_valid, dtype=tf.uint8, saturate=True)
outputs_node = tf.map_fn(tf.image.encode_png, converted_outputs, dtype=tf.string, name='output_pngs')
extracted_feature_gen = gen_output
extracted_feature_target = targets_batch
with tf.variable_scope('generator_loss'):
# Content loss
with tf.variable_scope('content_loss'):
diff = extracted_feature_gen - extracted_feature_target
content_loss = tf.reduce_mean(tf.reduce_sum(tf.square(diff), axis=[3]))
gen_loss = content_loss
with tf.variable_scope('generator_train'):
# Need to wait discriminator to perform train step
update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
with tf.control_dependencies(update_ops):
gen_tvars = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, scope='generator')
gen_optimizer = tf.train.AdamOptimizer(0.01)
gen_grads_and_vars = gen_optimizer.compute_gradients(gen_loss, gen_tvars)
gen_train = gen_optimizer.apply_gradients(gen_grads_and_vars, global_step=global_step)
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
saver = tf.train.Saver(max_to_keep=SAVER_MAX_KEEP)
sv = tf.train.Supervisor(logdir='./save', save_summaries_secs=0, saver=None)
step = 1
with sv.managed_session(config=config) as sess:
# Performing the training
for epoch_idx in range(0, EPOCH):
gen_dataset()
lr_paths = glob.glob('dataset/lr/*.jpg')
hr_paths = glob.glob('dataset/hr/*.jpg')
steps_per_epoch = len(lr_paths) // BATCH_SIZE
all_paths = list(zip(lr_paths, hr_paths))
lr_paths = [lr for lr, _ in all_paths]
hr_paths = [hr for _, hr in all_paths]
random.shuffle(all_paths)
epoch_step = 0
for idx in range(0, len(lr_paths) - BATCH_SIZE, BATCH_SIZE):
lr_batch = lr_paths[idx: idx + BATCH_SIZE]
hr_batch = hr_paths[idx: idx + BATCH_SIZE]
fetches = {
"global_step": global_step,
"train": gen_train,
"gen_output": gen_output,
"gen_loss": gen_loss
}
feed_dict = {
is_train: True,
lr_image_paths: lr_batch,
hr_image_paths: hr_batch,
}
results = sess.run(fetches, feed_dict=feed_dict)
if step % SHOW_INFO_INTERVAL == 0:
loss = results['gen_loss']
print('[%d][%d/%d] step:%d, loss:%.2f' % (epoch_idx, epoch_step, steps_per_epoch, step, loss))
if step % VALIDATE_INTERVAL == 0:
fetches = {
"valid_loss": valid_loss,
"outputs_node": outputs_node,
"gen_valid": gen_valid
}
feed_dict = {
is_train: False,
input_node: lr_valid_image,
target_node: hr_valid_image,
}
val_results = sess.run(fetches, feed_dict=feed_dict)
val_loss = val_results['valid_loss']
print('valid loss: %.2f' % val_loss)
save_images(val_results['outputs_node'][0],
filename='step_%d_loss_%.2f.jpg' % (step, val_loss))
if step % SAVE_MODEL_INTERVAL == 0:
print('saving ckpt.')
filename = f'step_{step}.ckpt'
saver.save(sess, f'model_out/{filename}')
epoch_step += 1
step += 1
