Update whitening and coloring transform

Author: tyui592

style_decorator.py

@@ -1,7 +1,7 @@
 import torch
 import torch.nn.functional as F
 
-from wct import batch_whitening, batch_coloring
+from wct import whitening, coloring
 
 def extract_patches(feature, patch_size, stride, padding='zero'):
     kh, kw = patch_size
@@ -136,17 +136,17 @@ def reassemble_feature(self, normalized_content_feature, normalized_style_featur
     def forward(self, content_feature, style_feature, style_strength=1.0, patch_size=3, patch_stride=1): 
 
         # 1-1. content feature projection
-        normalized_content_feature = batch_whitening(content_feature)
+        normalized_content_feature = whitening(content_feature)
 
         # 1-2. style feature projection
-        normalized_style_feature = batch_whitening(style_feature)
+        normalized_style_feature = whitening(style_feature)
 
         # 2. swap content and style features
         reassembled_feature = self.reassemble_feature(normalized_content_feature, normalized_style_feature,
                 patch_size=patch_size, patch_stride=patch_stride)
 
         # 3. reconstruction feature with style mean and covariance matrix
-        stylized_feature = batch_coloring(reassembled_feature, style_feature)
+        stylized_feature = coloring(reassembled_feature, style_feature)
 
         # 4. content and style interpolation
         result_feature = (1-style_strength) * content_feature + style_strength * stylized_feature

wct.py

@@ -1,61 +1,52 @@
 import torch
 
-def mean_covsqrt(f, inverse=False, eps=1e-5):
-    c, h, w = f.size()
-    
-    f_mean = torch.mean(f.view(c, h*w), dim=1, keepdim=True)
-    f_zeromean = f.view(c, h*w) - f_mean
-    f_cov = torch.mm(f_zeromean, f_zeromean.t())
-
-    u, s, v = torch.svd(f_cov)
-
-    k = c
-    for i in range(c):
-        if s[i] < eps:
-            k = i
-            break
-            
-    if inverse:
-        p = -0.5
-    else:
-        p = 0.5
-        
-    f_covsqrt = torch.mm(torch.mm(v[:, 0:k], torch.diag(s[0:k].pow(p))), v[:, 0:k].t())
-    return f_mean, f_covsqrt
+def covsqrt_mean(feature, inverse=False, tolerance=1e-14):
+    # I referenced the default svd tolerance value in matlab.
 
-def whitening(f):
-    c, h, w = f.size()
+    b, c, h, w = feature.size()
 
-    f_mean, f_inv_covsqrt = mean_covsqrt(f, inverse=True)
-    
-    whiten_f = torch.mm(f_inv_covsqrt, f.view(c, h*w) - f_mean)
-    
-    return whiten_f.view(c, h, w)
+    mean = torch.mean(feature.view(b, c, -1), dim=2, keepdim=True)
+    zeromean = feature.view(b, c, -1) - mean
+    cov = torch.bmm(zeromean, zeromean.transpose(1, 2))
 
-def coloring(f, t):
-    f_c, f_h, f_w = f.size()
-    t_c, t_h, t_w = t.size()
-    
-    t_mean, t_covsqrt = mean_covsqrt(t)
+    evals, evects = torch.symeig(cov, eigenvectors=True)
 
-    colored_f = torch.mm(t_covsqrt, f.view(f_c, f_h*f_w)) + t_mean
+    p = 0.5
+    if inverse:
+        p *= -1
+
+    covsqrt = []
+    for i in range(b):
+        k = 0
+        for j in range(c):
+            if evals[i][j] > tolerance:
+                k = j
+                break
+        covsqrt.append(torch.mm(evects[i][:, k:],
+                            torch.mm(evals[i][k:].pow(p).diag_embed(),
+                                     evects[i][:, k:].t())).unsqueeze(0))
+    covsqrt = torch.cat(covsqrt, dim=0)
+
+    u, s, v = torch.svd(cov)
+
+    return covsqrt, mean
 
-    return colored_f.view(f_c, f_h, f_w)
 
-def batch_whitening(f):
-    b, c, h, w = f.size()
+def whitening(feature):
+    b, c, h, w = feature.size()
+    
+    inv_covsqrt, mean = covsqrt_mean(feature, inverse=True)
 
-    whiten_f = torch.Tensor(b, c, h, w).type_as(f)
-    for i, f_ in enumerate(torch.split(f, 1)):
-        whiten_f[i] = whitening(f_.squeeze())
-        
-    return whiten_f
+    normalized_feature = torch.matmul(inv_covsqrt, feature.view(b, c, -1)-mean)
+    
+    return normalized_feature.view(b, c, h, w)
 
-def batch_coloring(f, t):
-    b, c, h, w = f.size()
 
-    colored_f = torch.Tensor(b, c, h, w).type_as(f)
-    for i, (f_, t_) in enumerate(zip(torch.split(f, 1), torch.split(t, 1))):
-        colored_f[i] = coloring(f_.squeeze(), t_.squeeze())
+def coloring(feature, target):
+    b, c, h, w = feature.size()
 
-    return colored_f
+    covsqrt, mean = covsqrt_mean(target)
+    
+    colored_feature = torch.matmul(covsqrt, feature.view(b, c, -1)) + mean
+    
+    return colored_feature.view(b, c, h, w)