Spaces:
Runtime error
Runtime error
File size: 1,912 Bytes
82ef366 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 |
def masked_adain(content_feat, style_feat, content_mask, style_mask):
assert (content_feat.size()[:2] == style_feat.size()[:2])
size = content_feat.size()
style_mean, style_std = calc_mean_std(style_feat, mask=style_mask)
content_mean, content_std = calc_mean_std(content_feat, mask=content_mask)
normalized_feat = (content_feat - content_mean.expand(size)) / content_std.expand(size)
style_normalized_feat = normalized_feat * style_std.expand(size) + style_mean.expand(size)
return content_feat * (1 - content_mask) + style_normalized_feat * content_mask
def calc_mean_std(feat, eps=1e-5, mask=None):
# eps is a small value added to the variance to avoid divide-by-zero.
size = feat.size()
if len(size) == 2:
return calc_mean_std_2d(feat, eps, mask)
assert (len(size) == 3)
C = size[0]
if mask is not None:
feat_var = feat.view(C, -1)[:, mask.view(-1) == 1].var(dim=1) + eps
feat_std = feat_var.sqrt().view(C, 1, 1)
feat_mean = feat.view(C, -1)[:, mask.view(-1) == 1].mean(dim=1).view(C, 1, 1)
else:
feat_var = feat.view(C, -1).var(dim=1) + eps
feat_std = feat_var.sqrt().view(C, 1, 1)
feat_mean = feat.view(C, -1).mean(dim=1).view(C, 1, 1)
return feat_mean, feat_std
def calc_mean_std_2d(feat, eps=1e-5, mask=None):
# eps is a small value added to the variance to avoid divide-by-zero.
size = feat.size()
assert (len(size) == 2)
C = size[0]
if mask is not None:
feat_var = feat.view(C, -1)[:, mask.view(-1) == 1].var(dim=1) + eps
feat_std = feat_var.sqrt().view(C, 1)
feat_mean = feat.view(C, -1)[:, mask.view(-1) == 1].mean(dim=1).view(C, 1)
else:
feat_var = feat.view(C, -1).var(dim=1) + eps
feat_std = feat_var.sqrt().view(C, 1)
feat_mean = feat.view(C, -1).mean(dim=1).view(C, 1)
return feat_mean, feat_std
|