def run_from_file_paths(model_args, input_args):
sess, generated_img, content_input, style_input = model_args
alpha, content_images, style_images, output_folder, progress_update = input_args
n_imgs = len(style_images) * len(content_images)
for i, content_image in enumerate(content_images):
content_image = get_img(content_image)
for j, style_image in enumerate(style_images):
style_image = get_img(style_image)
content_tensor = np.expand_dims(content_image, axis=0)
style_tensor = np.expand_dims(style_image, axis=0)
result = sess.run(
generated_img,
feed_dict={
content_input: content_tensor,
style_input: style_tensor
})
# assemble the output file name
style_img_split = os.path.splitext(os.path.basename(style_images[j]))
content_img_split = os.path.splitext(os.path.basename(content_images[i]))
output_image_path = content_img_split[0] + "-" + style_img_split[0] + content_img_split[1]
output_image_path = os.path.join(output_folder, output_image_path)
save_img(output_image_path, result[0])
# update the gimp progress bar
curr_img = float(i * len(style_images) + (j + 1))
progress_update(curr_img / n_imgs)
def run_from_file_paths(model_args, input_args):
sess, input_photo, output_photo = model_args
imgs, style_name, output_folder, progress_update, tmp = input_args
for i, content_path in enumerate(imgs):
img = get_img(content_path)
img_shape = img.shape[:2]
# Resize the smallest side of the image to the image_size
alpha = float(image_size) / float(min(img_shape))
img = imresize(img, size=alpha)
img = np.expand_dims(img, axis=0)
img = sess.run(output_photo,
feed_dict={
input_photo: normalize_arr_of_imgs(img),
})
img = denormalize_arr_of_imgs(img[0])
img = imresize(img, size=img_shape)
# generate output path
content_name = os.path.basename(content_path)
file_ext = os.path.splitext(content_name)
output_path = os.path.join(
output_folder, file_ext[0] + "-" + style_name + file_ext[1])
save_img(output_path, img)
progress_update((tmp["completed"] + float(i + 1)) / tmp["tot"])
def inference(from_file_path, args, batch_size=1, device_t='/gpu:0'):
imgs, style_name = args[0:2]
# check if the specified model exists
checkpoint_dir = os.path.join(pluginFolderPath, "models/",
style_name + ".ckpt")
assert os.path.exists(checkpoint_dir), 'Checkpoint not found!'
soft_config = tf.ConfigProto(allow_soft_placement=True)
soft_config.gpu_options.allow_growth = True
tf.reset_default_graph()
with tf.Session(config=soft_config) as sess:
img_shape = None
if from_file_path:
img_shape = get_img(imgs[0]).shape
else:
img_shape = imgs[0].shape
input_img_placeholder = tf.placeholder(dtype=tf.float32,
shape=(batch_size, ) +
img_shape,
name='input_img_placeholder')
preds = transform.net(input_img_placeholder)
saver = tf.train.Saver()
if os.path.isdir(checkpoint_dir):
ckpt = tf.train.get_checkpoint_state(checkpoint_dir)
if ckpt and ckpt.model_checkpoint_path:
saver.restore(sess, ckpt.model_checkpoint_path)
else:
raise Exception("No checkpoint found...")
else:
saver.restore(sess, checkpoint_dir)
model_args = (sess, input_img_placeholder, preds)
if from_file_path:
run_from_file_paths(model_args, args)
else:
return run_from_layers(model_args, args[0])
def run_from_file_paths(model_args, input_args):
sess, input_img_placeholder, preds = model_args
imgs, style_name, output_folder, progress_update, tmp = input_args
for i, content_path in enumerate(imgs):
img = get_img(content_path)
img = np.expand_dims(img, axis=0)
img = sess.run(preds, feed_dict={
input_img_placeholder: img,
})
img = np.clip(img[0], 0, 255).astype(np.uint8)
# generate output path
content_name = os.path.basename(content_path)
file_ext = os.path.splitext(content_name)
output_path = os.path.join(
output_folder, file_ext[0] + "-" + style_name + file_ext[1])
save_img(output_path, img)
progress_update((tmp["completed"] + float(i + 1)) / tmp["tot"])