def main():
args = parse_args()
if args is None:
exit()
content_image = utils.load_image(args.content, max_size=args.max_size)
style_image = utils.load_image(
args.style, shape=[content_image.shape[1], content_image.shape[0]])
init_image = None
if args.initial_type == 'content':
init_image = content_image
elif args.initial_type == 'style':
init_image = style_image
elif args.initial_type == 'random':
init_image = np.random.normal(size=content_image.shape,
scale=np.std(content_image))
assert init_image is not None
# create a map for content layers info
CONTENT_LAYERS = {}
for layer, weight in zip(args.content_layers, args.content_layer_weights):
CONTENT_LAYERS[layer] = weight
# create a map for style layers info
STYLE_LAYERS = {}
for layer, weight in zip(args.style_layers, args.style_layer_weights):
STYLE_LAYERS[layer] = weight
with tf.Session() as sess:
content_image = np.expand_dims(content_image, 0)
style_image = np.expand_dims(style_image, 0)
init_image = np.expand_dims(init_image, 0)
model = style_transfer.StyleTransfer(
style_transfer.Options(model_path=args.model_path,
content_image=content_image,
style_image=style_image,
init_image=init_image,
content_layers=CONTENT_LAYERS,
log_interval=args.log_interval,
image_log_dir=args.image_log_dir,
loss_ratio=args.loss_ratio,
steps=args.steps,
style_layers=STYLE_LAYERS))
final_image = model.generate(sess)
utils.save_image(final_image, args.output)
def _get_inputs(tfrecord_filename, style_image_files, image_size, batch_size):
# Get all of the images from the input dataset
dataset = dataset_builder.DatasetBuilder.build(tfrecord_filename,
batch_size, image_size)
dataset_iterator = dataset.make_one_shot_iterator()
# Load the style images.
logger.info('Loading style images:\n%s' % '\n'.join(style_image_files))
style_imgs = []
for filename in style_image_files:
# Note no preprocessing is done while loading.
img = utils.load_image(filename, *image_size)
style_imgs.append(img)
style_imgs = numpy.array(style_imgs)
return dataset_iterator, style_imgs
'--model-checkpoint',
type=str,
required=True,
help='Checkpoint from a trained Style Transfer Network.')
args = parser.parse_args()
logger.info('Loading model from %s' % args.model_checkpoint)
custom_objects = {
'InstanceNormalization':
keras_contrib.layers.normalization.InstanceNormalization,
'DeprocessStylizedImage': layers.DeprocessStylizedImage
}
transfer_net = keras.models.load_model(args.model_checkpoint,
custom_objects=custom_objects)
image_size = transfer_net.input_shape[1:3]
inputs = [transfer_net.input, keras.backend.learning_phase()]
outputs = [transfer_net.output]
transfer_style = keras.backend.function(inputs, outputs)
input_image = utils.load_image(args.input_image,
image_size[0],
image_size[1],
expand_dims=True)
output_image = transfer_style([input_image, 1])[0]
output_image = PIL.Image.fromarray(numpy.uint8(output_image[0]))
output_image.save(args.output_image)