
from __future__ import absolute_import

import sys
sys.path.append('..')
sys.path.append('.')
import PerceptualSimilarity as ps
import scipy
import scipy.misc
import numpy as np
import torch
from torch.autograd import Variable

use_gpu = False

ref_path  = './imgs/ex_ref.png'
pred_path = './imgs/ex_p1.png'

ref_img = scipy.misc.imread(ref_path).transpose(2, 0, 1) / 255.
pred_img = scipy.misc.imread(pred_path).transpose(2, 0, 1) / 255.

# Torchify
ref = Variable(torch.FloatTensor(ref_img), requires_grad=False)
pred = Variable(torch.FloatTensor(pred_img))
if(use_gpu):
    ref = ref.cuda()
    pred = pred.cuda()

# 1 x 3 x H x W
ref = ref.unsqueeze(0)
pred = pred.unsqueeze(0)

loss_fn = ps.PerceptualLoss(use_gpu=use_gpu)
dist = loss_fn.forward(pred, ref, normalize=True)

# As optimization, test backprop
class PerceptModel(torch.nn.Module):
    def __init__(self):
        super(PerceptModel, self).__init__()
        self.pred = torch.nn.Parameter(pred.data)

    def forward(self):
        return self.pred

model = PerceptModel()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3, betas=(0.9, 0.999))

import matplotlib.pyplot as plt
plt.ion()
fig = plt.figure(1)
ax = fig.add_subplot(131)
ax.imshow(ref_img.transpose(1, 2, 0))
ax.set_title('reference')
ax = fig.add_subplot(133)
ax.imshow(pred_img.transpose(1, 2, 0))
ax.set_title('orig pred')

for i in range(1000):
    optimizer.zero_grad()
    dist = loss_fn.forward(model.forward(), ref, normalize=True)
    dist.backward()
    optimizer.step()
    pred.data = torch.max(torch.clamp(pred.data, -1, 1))
    if i % 10 == 0:
        print('iter %d, dist %.3g' % (i, dist.view(-1).data.cpu().numpy()[0]))
        pred_img = model.pred[0].data.cpu().numpy().transpose(1, 2, 0)
        pred_img = np.clip(pred_img, 0, 1)
        ax = fig.add_subplot(132)            
        ax.imshow(pred_img)
        ax.set_title('iter %d, dist %.3g' % (i, dist.view(-1).data.cpu().numpy()[0]))
        plt.pause(5e-2)
        # plt.imsave('imgs_saved/%04d.jpg'%i,pred_img)



@@ -486,16 +486,29 @@
 ref_img)
+.unsqueeze(0)
 , requir
@@ -565,16 +565,49 @@
 red_img)
+.unsqueeze(0), requires_grad=True
 )%0Aif(use
@@ -662,429 +662,51 @@
 ()%0A%0A
-# 1 x 3 x H x W%0Aref = ref.unsqueeze(0)%0Apred = pred.unsqueeze(0)%0A%0Aloss_fn = ps.PerceptualLoss(use_gpu=use_gpu)%0Adist = loss_fn.forward(pred, ref, normalize=True)%0A%0A# As optimization, test backprop%0Aclass PerceptModel(torch.nn.Module):%0A    def __init__(self):%0A        super(PerceptModel, self).__init__()%0A        self.pred = torch.nn.Parameter(pred.data)%0A%0A    def forward(self):%0A        return self.pred%0A%0Amodel = PerceptModel(
+loss_fn = ps.PerceptualLoss(use_gpu=use_gpu
 )%0Aop
@@ -736,26 +736,15 @@
 dam(
-model.parameters()
+%5Bpred,%5D
 , lr
@@ -1041,34 +1041,8 @@
 0):%0A
-    optimizer.zero_grad()%0A
     
@@ -1068,23 +1068,12 @@
 ard(
-model.forward()
+pred
 , re
@@ -1091,16 +1091,42 @@
 e=True)%0A
+    optimizer.zero_grad()%0A
     dist
@@ -1178,18 +1178,8 @@
 a = 
-torch.max(
 torc
@@ -1201,15 +1201,18 @@
 ta, 
--1
+0
 , 1)
-)
+%0A    
 %0A   
@@ -1330,14 +1330,8 @@
 g = 
-model.
 pred

clean up perceptual loss code