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 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)
for loss in losses: loss = loss.cuda() optImg.cuda() ########### TRAINING ########## for iteration in range(1, opt.niter + 1): print('Iteration [%d]/[%d]' % (iteration, opt.niter)) def closure(): optimizer.zero_grad() out = vgg(optImg, loss_layers) totalLossList = [] for i in range(len(out)): layer_output = out[i] loss_i = losses[i] target_i = targets[i] totalLossList.append(loss_i(layer_output, target_i) * weights[i]) totalLoss = sum(totalLossList) totalLoss.backward() print('loss: %f' % (totalLoss.data[0])) return totalLoss optimizer.step(closure) outImg = optImg.data[0].cpu() if (opt.luminance_only): outImg = np.expand_dims(outImg.numpy(), 0) outImg = util.join_yiq_to_bgr(outImg, content_iq) save_image(torch.from_numpy(outImg).squeeze()) else: save_image(outImg.squeeze())
# shift everything to cuda if possible if(opt.cuda): for loss in losses: loss = loss.cuda() optImg.cuda() ########### TRAINING ########## for iteration in range(1,opt.niter+1): print('Iteration [%d]/[%d]'%(iteration,opt.niter)) def closure(): optimizer.zero_grad() out = vgg(optImg,loss_layers) totalLossList = [] for i in range(len(out)): layer_output = out[i] loss_i = losses[i] target_i = targets[i] totalLossList.append(loss_i(layer_output,target_i) * weights[i]) totalLoss = sum(totalLossList) totalLoss.backward() print('loss: %f'%(totalLoss.data[0])) return totalLoss optimizer.step(closure) outImg = optImg.data[0].cpu() if(opt.luminance_only): outImg = np.expand_dims(outImg.numpy(),0) outImg = util.join_yiq_to_bgr(outImg,content_iq) save_image(torch.from_numpy(outImg).squeeze()) else: save_image(outImg.squeeze())