|
- import numpy as np
- import paddle
- import matplotlib.pyplot as plt
- import argparse
- import paddle_lpips as lpips
-
- parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
- parser.add_argument('--ref_path', type=str, default='./imgs/ex_ref.png')
- parser.add_argument('--pred_path', type=str, default='./imgs/ex_p1.png')
- parser.add_argument('--use_gpu', action='store_true', help='turn on flag to use GPU')
-
- opt = parser.parse_args()
-
- if(opt.use_gpu):
- paddle.set_device('gpu')
- else:
- paddle.set_device('cpu')
-
- loss_fn = lpips.LPIPS(net='vgg')
-
- ref = lpips.im2tensor(lpips.load_image(opt.ref_path))
- pred = lpips.im2tensor(lpips.load_image(opt.pred_path))
- pred.stop_gradient = False
-
- optimizer = paddle.optimizer.Adam(parameters=[pred,], learning_rate=1e-3, beta1=0.9, beta2=0.999)
-
- plt.ion()
- fig = plt.figure(1)
- ax = fig.add_subplot(131)
- ax.imshow(lpips.tensor2im(ref))
- ax.set_title('target')
- ax = fig.add_subplot(133)
- ax.imshow(lpips.tensor2im(pred))
- ax.set_title('initialization')
-
- for i in range(1000):
- dist = loss_fn.forward(pred, ref)
- optimizer.zero_grad()
- dist.backward()
- optimizer.step()
- pred[:] = paddle.clip(pred, -1, 1)
-
- if i % 10 == 0:
- print('iter %d, dist %.3g' % (i, dist.reshape((-1,)).numpy()[0]))
- pred[:] = paddle.clip(pred, -1, 1)
- pred_img = lpips.tensor2im(pred)
-
- ax = fig.add_subplot(132)
- ax.imshow(pred_img)
- ax.set_title('iter %d, dist %.3f' % (i, dist.reshape((-1,)).numpy()[0]))
- plt.pause(5e-2)
- # plt.imsave('imgs_saved/%04d.jpg'%i,pred_img)
-
|