Spaces:
Runtime error
Runtime error
| import torch | |
| from skimage.exposure import match_histograms | |
| import numpy as np | |
| 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)), | |
| multichannel=True) | |
| 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, target): | |
| # 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) | |
| 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) | |
| 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()) | |