def run(args):
if args.out_dir != None:
if not os.path.exists(args.out_dir):
try:
os.mkdir(args.out_dir)
except:
print 'cannot make directory {}'.format(args.out_dir)
exit()
elif not os.path.isdir(args.out_dir):
print 'file path {} exists but is not directory'.format(args.out_dir)
exit()
if args.type == 'vgg19':
vgg = VGG19()
else:
vgg = VGG()
content_image = open_and_resize_image(args.content, args.width, vgg)
print 'loading content image completed'
style_image = open_and_resize_image(args.style, args.width, vgg)
if args.match_color_histogram:
style_image = util.match_color_histogram(style_image, content_image)
if args.luminance_only:
content_image, content_iq = util.split_bgr_to_yiq(content_image)
style_image, style_iq = util.split_bgr_to_yiq(style_image)
content_mean = np.mean(content_image, axis=(1,2,3), keepdims=True)
content_std = np.std(content_image, axis=(1,2,3), keepdims=True)
style_mean = np.mean(style_image, axis=(1,2,3), keepdims=True)
style_std = np.std(style_image, axis=(1,2,3), keepdims=True)
style_image = (style_image - style_mean) / style_std * content_std + content_mean
print 'loading style image completed'
serializers.load_hdf5(args.model, vgg)
print 'loading neural network model completed'
optimizer = LBFGS(args.lr)
content_layers = args.content_layers.split(',')
style_layers = args.style_layers.split(',')
def on_epoch_done(epoch, x, losses):
if (epoch + 1) % args.save_iter == 0:
image = cuda.to_cpu(x.data)
if args.luminance_only:
image = util.join_yiq_to_bgr(image, content_iq)
image = vgg.postprocess(image[0], output_type='RGB').clip(0, 255).astype(np.uint8)
Image.fromarray(image).save(os.path.join(args.out_dir, 'out_{0:04d}.png'.format(epoch + 1)))
print 'epoch {} done'.format(epoch + 1)
print 'losses:'
label_width = max(map(lambda (name, loss): len(name), losses))
for name, loss in losses:
print ' {0:{width}s}: {1:f}'.format(name, loss, width=label_width)
if args.method == 'mrf':
model = MRF(vgg, optimizer, args.content_weight, args.style_weight, args.tv_weight, content_layers, style_layers, args.resolution_num, args.gpu, initial_image=args.initial_image, keep_color=args.keep_color)
else:
model = NeuralStyle(vgg, optimizer, args.content_weight, args.style_weight, args.tv_weight, content_layers, style_layers, args.resolution_num, args.gpu, initial_image=args.initial_image, keep_color=args.keep_color)
out_image = model.fit(content_image, style_image, args.iter, on_epoch_done)
out_image = cuda.to_cpu(out_image.data)
if args.luminance_only:
out_image = util.join_yiq_to_bgr(out_image, content_iq)
image = vgg.postprocess(out_image[0], output_type='RGB').clip(0, 255).astype(np.uint8)
Image.fromarray(image).save(os.path.join(args.out_dir, 'out.png'))
def style_transfer():
enc = VGG19(input_shape)
decs = [Decoder(enc, i, name='decoder %d' % i) for i in indices]
sty_net = StyleTransfer(enc, decs)
X = T.tensor4('input')
Y = T.tensor4('style')
mon = nn.Monitor(current_folder=args.path_to_weight_files)
for idx, layer in enumerate(indices):
weights = mon.load('decoder-%d-final.npz' % layer)
nn.utils.numpy2shared(weights, decs[idx].params)
nn.set_training_off()
X_styled = sty_net(X, Y)
transfer = nn.function([X, Y], X_styled, name='transfer style')
if os.path.isfile(input_path_test) and os.path.isfile(style_path_test):
input = prep_image_test(misc.imread(input_path_test))
style = prep_image_test(misc.imread(style_path_test))
output = transfer(input, style)
mon.imwrite('test %s' % input_path_test[:-4], input)
mon.imwrite('test %s' % style_path_test[:-4], style)
mon.imwrite('test %s %s' % (input_path_test[:-4], style_path_test[:-4]), output)
elif os.path.isfile(input_path_test) and os.path.isdir(style_path_test):
input = prep_image_test(misc.imread(input_path_test))
style_files = os.listdir(style_path_test)
for style_file in style_files:
style = prep_image_test(misc.imread(os.path.join(style_path_test, style_file)))
output = transfer(input, style)
mon.imwrite('test %s' % style_file[:-4], style)
mon.imwrite('test %s %s' % (input_path_test[:-4], style_file[:-4]), output)
mon.imwrite('test %s' % input_path_test[:-4], input)
elif os.path.isdir(input_path_test) and os.path.isfile(style_path_test):
style = prep_image_test(misc.imread(style_path_test))
input_files = os.listdir(input_path_test)
for input_file in input_files:
input = prep_image_test(misc.imread(os.path.join(input_path_test, input_file)))
output = transfer(input, style)
mon.imwrite('test %s' % input_file[:-4], input)
mon.imwrite('test %s %s' % (input_file[:-4], style_path_test[:-4]), output)
mon.imwrite('test %s' % style_path_test[:-4], style)
else:
style_files = os.listdir(style_path_test)
input_files = os.listdir(input_path_test)
for style_file in style_files:
style = prep_image_test(misc.imread(os.path.join(style_path_test, style_file)))
for input_file in input_files:
input = prep_image_test(misc.imread(os.path.join(input_path_test, input_file)))
output = transfer(input, style)
mon.imwrite('test %s' % input_file[:-4], input)
mon.imwrite('test %s %s' % (input_file[:-4], style_file[:-4]), output)
mon.imwrite('test %s' % style_file[:-4], style)
mon.flush()
print('Testing finished!')
def train():
enc = VGG19(input_shape)
decs = [Decoder(enc, i, name='decoder %d' % i) for i in indices]
sty_net = StyleTransfer(enc, decs)
X = T.tensor4('input')
Y = T.tensor4('style')
idx = T.scalar('iter', 'int32')
X_ = nn.placeholder((bs,) + input_shape[1:], name='input_plhd')
Y_ = nn.placeholder((bs,) + input_shape[1:], name='style_plhd')
lr_ = nn.placeholder(value=lr, name='lr_plhd')
nn.set_training_on()
losses = [dec.cost(X) for dec in decs]
updates = [nn.adam(loss[0] + weight * loss[1], dec.trainable, lr) for loss, dec in zip(losses, decs)]
nn.anneal_learning_rate(lr_, idx, 'inverse', decay=decay)
trains = [nn.function([], [loss[0], loss[1], dec(X, True)], givens={X: X_}, updates=update, name='train decoder')
for loss, dec, update in zip(losses, decs, updates)]
nn.set_training_off()
X_styled = sty_net(X, Y)
transfer = nn.function([], X_styled, givens={X: X_, Y: Y_}, name='transfer style')
data_train = DataManager(X_, input_path_train, bs, n_epochs, True, num_val_imgs=num_val_imgs, input_shape=input_shape)
data_test = DataManagerStyleTransfer((X_, Y_), (input_path_val, style_path_val), bs, 1, input_shape=input_shape)
mon = nn.Monitor(model_name='WCT', valid_freq=print_freq)
print('Training...')
for it in data_train:
results = [train(it) for train in trains]
with mon:
for layer, res in zip(indices, results):
if np.isnan(res[0] + res[1]) or np.isinf(res[0] + res[1]):
raise ValueError('Training failed!')
mon.plot('pixel loss at layer %d' % layer, res[0])
mon.plot('feature loss at layer %d' % layer, res[1])
if it % val_freq == 0:
mon.imwrite('recon img at layer %d' % layer, res[2])
for i in data_test:
img = transfer()
mon.imwrite('stylized image %d' % i, img)
mon.imwrite('input %d' % i, X_.get_value())
mon.imwrite('style %d' % i, Y_.get_value())
for idx, dec in zip(indices, decs):
mon.dump(nn.utils.shared2numpy(dec.params), 'decoder-%d.npz' % idx, 5)
mon.flush()
for idx, dec in zip(indices, decs):
mon.dump(nn.utils.shared2numpy(dec.params), 'decoder-%d-final.npz' % idx)
print('Training finished!')
break
if a[1].data.shape != b[1].data.shape:
match = False
break
if not match:
print 'Ignore %s because of parameter mismatch' % child.name
continue
for a, b in zip(child.namedparams(), dst_child.namedparams()):
b[1].data = a[1].data
print 'Copy %s' % child.name
if len(sys.argv) > 1 and sys.argv[1] == 'vgg19':
print 'load VGG19 caffemodel'
ref = CaffeFunction('VGG_ILSVRC_19_layers.caffemodel')
vgg = VGG19()
print 'copy weights'
copy_model(ref, vgg)
print 'save "vgg19.model"'
serializers.save_hdf5('vgg19.model', vgg)
else:
print 'load VGG16 caffemodel'
ref = CaffeFunction('VGG_ILSVRC_16_layers.caffemodel')
vgg = VGG()
print 'copy weights'
copy_model(ref, vgg)
print 'save "vgg16.model"'
serializers.save_hdf5('vgg16.model', vgg)
def preprocess_image(image, image_size, clip_rect=None):
if clip_rect is not None:
image = image.crop(clip_rect)
image = image.resize(image_size, Image.BILINEAR)
x = np.asarray(image, dtype=np.float32)
return VGG19.preprocess(x, input_type='RGB')
def train(args,
image_path,
source_image_paths,
target_image_paths,
input_clip_rect=None,
clip_rect=None):
iteration = args.iter
batch_size = args.batch_size
device_id = args.gpu
lr = args.lr
tv_weight = args.tv_weight
near_image_num = args.near_image
make_dir(os.path.split(args.output_image)[0])
net = VGG19()
serializers.load_npz(args.model, net)
if device_id >= 0:
net.to_gpu(device_id)
xp = net.xp
original_image = Image.open(image_path).convert('RGB')
if input_clip_rect is not None:
original_image = original_image.crop(input_clip_rect)
image = preprocess_image(original_image, input_image_size)
image_mean = np.mean(image, axis=(2, 3), keepdims=True)
image_std = np.std(image, axis=(2, 3), keepdims=True)
x = xp.asarray(image)
org_layers = feature(net, x)
org_layers = [layer.data for layer in org_layers]
org_layer_norms = [
xp.asarray(np.linalg.norm(cuda.to_cpu(layer), axis=1, keepdims=True))
for layer in org_layers
]
print('Calculating source feature')
if len(source_image_paths) > near_image_num:
source_image_paths = rank_image(net, source_image_paths,
input_image_size, image,
near_image_num, clip_rect)
source_feature = mean_feature(net, source_image_paths, input_image_size,
org_layers[-1], near_image_num, batch_size,
clip_rect)
print('Calculating target feature')
if len(target_image_paths) > near_image_num:
target_image_paths = rank_image(net, target_image_paths,
input_image_size, image,
near_image_num, clip_rect)
target_feature = mean_feature(net, target_image_paths, input_image_size,
org_layers[-1], near_image_num, batch_size,
clip_rect)
attribute_vectors = [
normalized_diff(s, t) for s, t in zip(source_feature, target_feature)
]
base, ext = os.path.splitext(args.output_image)
residuals = []
initial_x = xp.random.uniform(-10, 10, x.shape).astype(np.float32)
print('Calculating residuals')
link = chainer.Link(x=x.shape)
if device_id >= 0:
link.to_gpu(device_id)
link.x.data[...] = initial_x
optimizer = LBFGS(lr, stack_size=5)
optimizer.setup(link)
for j in six.moves.range(600):
losses = update(net, optimizer, link, org_layers, tv_weight)
if (j + 1) % 20 == 0:
z = cuda.to_cpu(link.x.data)
z = adjust_color_distribution(z, image_mean, image_std)
residuals.append(z - image)
for i in six.moves.range(1, 6):
w = i * 0.1
print('Generating image for weight: {0:.2f}'.format(w))
link = chainer.Link(x=x.shape)
if device_id >= 0:
link.to_gpu(device_id)
link.x.data[...] = initial_x
target_layers = [
layer + w * n * a for layer, n, a in zip(
org_layers, org_layer_norms, attribute_vectors)
]
optimizer = LBFGS(lr, stack_size=5)
optimizer.setup(link)
for j in six.moves.range(iteration):
losses = update(net, optimizer, link, target_layers, tv_weight)
if (j + 1) % 100 == 0:
print('iter {} done loss:'.format(j + 1))
print(losses)
if (j + 1) % 500 == 0:
z = cuda.to_cpu(link.x.data)
z = adjust_color_distribution(z, image_mean, image_std)
z -= find_nearest(residuals, z - image)
file_name = '{0}_{1:02d}_{2:04d}{3}'.format(
base, i, j + 1, ext)
postprocess_image(original_image, z - image).save(file_name)
z = cuda.to_cpu(link.x.data)
z = adjust_color_distribution(z, image_mean, image_std)
z -= find_nearest(residuals, z - image)
file_name = '{0}_{1:02d}{2}'.format(base, i, ext)
postprocess_image(original_image, z - image).save(file_name)
print('Completed')
def save_image(path, x):
x = VGG19.postprocess(x, output_type='RGB')
x = x.clip(0, 255).astype(np.uint8)
Image.fromarray(x).save(path)
def load_image(path, image_size):
image = Image.open(path).resize(image_size)
x = np.asarray(image, dtype=np.float32)
return VGG19.preprocess(x, input_type='RGB')
def main():
args = parse_arg()
iteration = args.iter
batch_size = args.batch_size
device_id = args.gpu
lr = args.lr
tv_weight = args.tv_weight
max_image = args.max_image
near_image_num = args.near_image
net = VGG19()
serializers.load_npz(args.model, net)
if device_id >= 0:
net.to_gpu(device_id)
xp = net.xp
image = load_image(args.input_image, input_image_size)
x = xp.asarray(image)
org_layers = feature(net, x)
org_layers = [layer.data for layer in org_layers]
print('Calculating source feature')
source_image_files = list_dir_image(args.source_dir, max_image)
if len(source_image_files) > near_image_num:
source_image_files = rank_image(net, source_image_files,
input_image_size, image,
near_image_num)
source_feature = mean_feature(net, source_image_files, input_image_size,
org_layers[-1], near_image_num, batch_size)
print('Calculating target feature')
target_image_files = list_dir_image(args.target_dir, max_image)
if len(target_image_files) > near_image_num:
target_image_files = rank_image(net, target_image_files,
input_image_size, image,
near_image_num)
target_feature = mean_feature(net, target_image_files, input_image_size,
org_layers[-1], near_image_num, batch_size)
attribute_vector = [
feature_diff(s, t) for s, t in zip(source_feature, target_feature)
]
base, ext = os.path.splitext(args.output_image)
for i in six.moves.range(1, 11):
w = i * 0.2
print('Generating image for weight: {0:.2f}'.format(w))
link = chainer.Link(x=x.shape)
if device_id >= 0:
link.to_gpu(device_id)
link.x.data[...] = xp.random.uniform(-10, 10,
x.shape).astype(np.float32)
target_layers = [
layer + w * a for layer, a in zip(org_layers, attribute_vector)
]
optimizer = LBFGS(lr, size=5)
optimizer.setup(link)
for j in six.moves.range(iteration):
losses = update(net, optimizer, link, target_layers, tv_weight)
print(losses)
if (j + 1) % 500 == 0:
print('iter {} done loss:'.format(j + 1))
save_image(
'{0}_{1:02d}_{2:04d}{3}'.format(base, i, j + 1, ext),
cuda.to_cpu(link.x.data))
save_image('{0}_{1:02d}{2}'.format(base, i, ext),
cuda.to_cpu(link.x.data))
print('Completed')
def preprocess_image(image, image_size):
image = image.resize(image_size, Image.BILINEAR)
x = np.asarray(image, dtype=np.float32)
return VGG19.preprocess(x, input_type='RGB')