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