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 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 optimize(net, link, target_layers, iteration):
optimizer = LBFGS(size=10)
optimizer.setup(link)
for i in six.moves.range(iteration):
update(net, optimizer, link, target_layers)
return link.x.data
