implement Lu et al transfer

Author: black-puppydog

WCT.py

@@ -3,6 +3,7 @@
 import os
 import torch
 import argparse
+import pprint
 from PIL import Image
 from torch.autograd import Variable
 import torchvision.utils as vutils
@@ -23,6 +24,9 @@
 parser.add_argument('--decoder2', default='models/vgg19_normalized_decoder2.pth.tar', help='Path to the decoder2')
 parser.add_argument('--decoder1', default='models/vgg19_normalized_decoder1.pth.tar', help='Path to the decoder1')
 parser.add_argument('--cuda', action='store_true', help='enables cuda')
+parser.add_argument('--transform-method', choices=['original', 'closed-form'], default='original',
+                    help=('How to whiten and color the features. "original" for the formulation of Li et al. ( https://arxiv.org/abs/1705.08086 )  '
+                          'or "closed-form" for method of Lu et al. ( https://arxiv.org/abs/1906.00668 '))
 parser.add_argument('--batch_size', type=int, default=1, help='batch size')
 parser.add_argument('--fineSize', type=int, default=512, help='resize image to fineSize x fineSize,leave it to 0 if not resize')
 parser.add_argument('--outf', default='samples/', help='folder to output images')
@@ -32,6 +36,7 @@
 parser.add_argument('--gpu', type=int, default=0, help="which gpu to run on.  default is 0")
 
 args = parser.parse_args()
+pprint.pprint(args.__dict__, indent=2)
 
 try:
     os.makedirs(args.outf)
@@ -44,7 +49,7 @@
                                      batch_size=1,
                                      shuffle=False)
 
-def styleTransfer(wct, targets, contentImg, styleImg, imname, gamma, delta, outf):
+def styleTransfer(wct, targets, contentImg, styleImg, imname, gamma, delta, outf, transform_method):
 
   current_result = contentImg
   eIorigs = [f.cpu().squeeze(0) for f in wct.encoder(contentImg, targets)]
@@ -58,8 +63,8 @@ def styleTransfer(wct, targets, contentImg, styleImg, imname, gamma, delta, outf
     else:
       eIlast = wct.encoder(current_result, target).cpu().squeeze(0)
 
-    CsIlast = wct.transform(eIlast, eIs).float()
-    CsIorig = wct.transform(eIorig, eIs).float()
+    CsIlast = wct.transform(eIlast, eIs, transform_method).float()
+    CsIorig = wct.transform(eIorig, eIs, transform_method).float()
 
     decoder_input = (gamma*(delta * CsIlast + (1-delta) * CsIorig) \
                      + (1-gamma) * eIorig)
@@ -91,7 +96,7 @@ def main():
       # WCT Style Transfer
       targets = [f'relu{t}_1' for t in args.targets]
       styleTransfer(wct, targets, contentImg, styleImg, imname,
-                    args.gamma, args.delta, args.outf)
+                    args.gamma, args.delta, args.outf, args.transform_method)
       end_time = time.time()
       print(' Elapsed time is: %f' % (end_time - start_time))
       avgTime += (end_time - start_time)

util.py

@@ -12,6 +12,10 @@
 
 
 def matrix_sqrt(A):
+  A = A.clone()
+  a_diag_ = A.diagonal()
+  a_diag_ += 1e-4
+
   s_u, s_e, s_v = torch.svd(A,some=False)
 
   k_s = A.shape[-1]
@@ -27,6 +31,9 @@ def matrix_sqrt(A):
 
 
 def matrix_inv_sqrt(A):
+  A = A.clone()
+  a_diag_ = A.diagonal()
+  a_diag_ += 1e-4
   k_c = A.shape[-1]
   c_u,c_e,c_v = torch.svd(A, some=False)
 
@@ -65,36 +72,49 @@ def __init__(self,args):
                          'relu4_1': self.d4,
                          'relu5_1': self.d5}
 
-    def whiten_and_color(self,cF,sF):
+    def whiten_and_color(self,cF,sF, method):
         cFSize = cF.size()
+        print(f'cF.shape = {cF.shape}')
         c_mean = torch.mean(cF,1) # c x (h x w)
         c_mean = c_mean.unsqueeze(1).expand_as(cF)
         cF = cF - c_mean
 
         contentConv = torch.mm(cF,cF.t()).div(cFSize[1]-1) + torch.eye(cFSize[0]).double()
-        cF_inv_sqrt = matrix_inv_sqrt(contentConv)
 
         sFSize = sF.size()
         s_mean = torch.mean(sF,1)
         sF = sF - s_mean.unsqueeze(1).expand_as(sF)
         styleConv = torch.mm(sF,sF.t()).div(sFSize[1]-1)
-        sF_sqrt = matrix_sqrt(styleConv)
 
-        whiten_cF = torch.mm(cF_inv_sqrt, cF)
+        if method == 'original':  # the original WCT by Li et al.
+          cF_inv_sqrt = matrix_inv_sqrt(contentConv)
+          sF_sqrt = matrix_sqrt(styleConv)
+          # whiten_cF = torch.mm(cF_inv_sqrt, cF)
+          # targetFeature = torch.mm(sF_sqrt,whiten_cF)
+          targetFeature = sF_sqrt @ (cF_inv_sqrt @ cF)
+        else:  # Lu et al.
+          assert method == 'closed-form'
+          cF_sqrt = matrix_sqrt(contentConv)
+          cF_inv_sqrt = matrix_inv_sqrt(contentConv)
+          print(f'cF_sqrt.shape = {cF_sqrt.shape}')
+          middle_matrix = matrix_sqrt(cF_sqrt @ styleConv @ cF_sqrt)
+          print(f'middle_matrix.shape = {middle_matrix.shape}')
+          transform_matrix = cF_inv_sqrt @ middle_matrix @ cF_inv_sqrt
+          targetFeature = transform_matrix @ cF
+          print(f'targetFeature.shape = {targetFeature.shape}')
 
-        targetFeature = torch.mm(sF_sqrt,whiten_cF)
         targetFeature = targetFeature + s_mean.unsqueeze(1).expand_as(targetFeature)
         return targetFeature
 
-    def transform(self, cF, sF):
+    def transform(self, cF, sF, method):
         cF = cF.double()
         sF = sF.double()
         C,W,H = cF.size(0),cF.size(1),cF.size(2)
         _,W1,H1 = sF.size(0),sF.size(1),sF.size(2)
         cFView = cF.view(C,-1)
         sFView = sF.view(C,-1)
 
-        targetFeature = self.whiten_and_color(cFView,sFView)
+        targetFeature = self.whiten_and_color(cFView, sFView, method)
         targetFeature = targetFeature.view_as(cF)
         return targetFeature