141 lines
4.6 KiB
Python
141 lines
4.6 KiB
Python
import numpy as np
|
|
import torch
|
|
from skimage.exposure import match_histograms
|
|
|
|
|
|
def calc_mean_std(feat, eps=1e-5):
|
|
# eps is a small value added to the variance to avoid divide-by-zero.
|
|
size = feat.size()
|
|
assert len(size) == 4
|
|
N, C = size[:2]
|
|
feat_var = feat.view(N, C, -1).var(dim=2) + eps
|
|
feat_std = feat_var.sqrt().view(N, C, 1, 1)
|
|
feat_mean = feat.view(N, C, -1).mean(dim=2).view(N, C, 1, 1)
|
|
return feat_mean, feat_std
|
|
|
|
|
|
def adaptive_instance_normalization(content_feat, style_feat):
|
|
assert content_feat.size()[:2] == style_feat.size()[:2]
|
|
size = content_feat.size()
|
|
style_mean, style_std = calc_mean_std(style_feat)
|
|
content_mean, content_std = calc_mean_std(content_feat)
|
|
|
|
normalized_feat = (content_feat - content_mean.expand(size)) / content_std.expand(
|
|
size
|
|
)
|
|
return normalized_feat * style_std.expand(size) + style_mean.expand(size)
|
|
|
|
|
|
## AdaMean
|
|
def adaptive_mean_normalization(content_feat, style_feat):
|
|
assert content_feat.size()[:2] == style_feat.size()[:2]
|
|
size = content_feat.size()
|
|
style_mean, style_std = calc_mean_std(style_feat)
|
|
content_mean, content_std = calc_mean_std(content_feat)
|
|
|
|
normalized_feat = content_feat - content_mean.expand(size)
|
|
return normalized_feat + style_mean.expand(size)
|
|
|
|
|
|
## AdaStd
|
|
def adaptive_std_normalization(content_feat, style_feat):
|
|
assert content_feat.size()[:2] == style_feat.size()[:2]
|
|
size = content_feat.size()
|
|
style_mean, style_std = calc_mean_std(style_feat)
|
|
content_mean, content_std = calc_mean_std(content_feat)
|
|
|
|
normalized_feat = (content_feat) / content_std.expand(size)
|
|
return normalized_feat * style_std.expand(size)
|
|
|
|
|
|
## EFDM
|
|
def exact_feature_distribution_matching(content_feat, style_feat):
|
|
assert content_feat.size() == style_feat.size()
|
|
B, C, W, H = (
|
|
content_feat.size(0),
|
|
content_feat.size(1),
|
|
content_feat.size(2),
|
|
content_feat.size(3),
|
|
)
|
|
value_content, index_content = torch.sort(
|
|
content_feat.view(B, C, -1)
|
|
) # sort conduct a deep copy here.
|
|
value_style, _ = torch.sort(
|
|
style_feat.view(B, C, -1)
|
|
) # sort conduct a deep copy here.
|
|
inverse_index = index_content.argsort(-1)
|
|
new_content = content_feat.view(B, C, -1) + (
|
|
value_style.gather(-1, inverse_index) - content_feat.view(B, C, -1).detach()
|
|
)
|
|
|
|
return new_content.view(B, C, W, H)
|
|
|
|
|
|
## HM
|
|
def histogram_matching(content_feat, style_feat):
|
|
assert content_feat.size() == style_feat.size()
|
|
B, C, W, H = (
|
|
content_feat.size(0),
|
|
content_feat.size(1),
|
|
content_feat.size(2),
|
|
content_feat.size(3),
|
|
)
|
|
x_view = content_feat.view(-1, W, H)
|
|
image1_temp = match_histograms(
|
|
np.array(x_view.detach().clone().cpu().float().transpose(0, 2)),
|
|
np.array(
|
|
style_feat.view(-1, W, H).detach().clone().cpu().float().transpose(0, 2)
|
|
),
|
|
)
|
|
image1_temp = (
|
|
torch.from_numpy(image1_temp)
|
|
.float()
|
|
.to(content_feat.device)
|
|
.transpose(0, 2)
|
|
.view(B, C, W, H)
|
|
)
|
|
return content_feat + (image1_temp - content_feat).detach()
|
|
|
|
|
|
def _calc_feat_flatten_mean_std(feat):
|
|
# takes 3D feat (C, H, W), return mean and std of array within channels
|
|
assert feat.size()[0] == 3
|
|
assert isinstance(feat, torch.FloatTensor)
|
|
feat_flatten = feat.view(3, -1)
|
|
mean = feat_flatten.mean(dim=-1, keepdim=True)
|
|
std = feat_flatten.std(dim=-1, keepdim=True)
|
|
return feat_flatten, mean, std
|
|
|
|
|
|
def _mat_sqrt(x):
|
|
U, D, V = torch.svd(x)
|
|
return torch.mm(torch.mm(U, D.pow(0.5).diag()), V.t())
|
|
|
|
|
|
def coral(source: torch.Tensor, target: torch.Tensor):
|
|
# assume both source and target are 3D array (C, H, W)
|
|
# Note: flatten -> f
|
|
|
|
source_f, source_f_mean, source_f_std = _calc_feat_flatten_mean_std(source.to('cpu'))
|
|
source_f_norm = (
|
|
source_f - source_f_mean.expand_as(source_f)
|
|
) / source_f_std.expand_as(source_f)
|
|
source_f_cov_eye = torch.mm(source_f_norm, source_f_norm.t()) + torch.eye(3)
|
|
|
|
target_f, target_f_mean, target_f_std = _calc_feat_flatten_mean_std(target.to('cpu'))
|
|
target_f_norm = (
|
|
target_f - target_f_mean.expand_as(target_f)
|
|
) / target_f_std.expand_as(target_f)
|
|
target_f_cov_eye = torch.mm(target_f_norm, target_f_norm.t()) + torch.eye(3)
|
|
|
|
source_f_norm_transfer = torch.mm(
|
|
_mat_sqrt(target_f_cov_eye),
|
|
torch.mm(torch.inverse(_mat_sqrt(source_f_cov_eye)), source_f_norm),
|
|
)
|
|
|
|
source_f_transfer = source_f_norm_transfer * target_f_std.expand_as(
|
|
source_f_norm
|
|
) + target_f_mean.expand_as(source_f_norm)
|
|
|
|
return source_f_transfer.view(source.size())
|