Spaces:
Running
Running
| # -*- coding: utf-8 -*- | |
| # Copyright (c) XiMing Xing. All rights reserved. | |
| # Author: XiMing Xing | |
| # Description: | |
| import math | |
| import torch | |
| import torch.nn as nn | |
| import torchvision | |
| import numpy as np | |
| class VGG16Extractor(nn.Module): | |
| def __init__(self, space): | |
| super().__init__() | |
| # load pretrained model | |
| self.vgg_layers = torchvision.models.vgg16( | |
| weights=torchvision.models.VGG16_Weights.DEFAULT | |
| ).features | |
| for param in self.parameters(): | |
| param.requires_grad = False | |
| self.capture_layers = [1, 3, 6, 8, 11, 13, 15, 22, 29] | |
| self.space = space | |
| def forward_base(self, x): | |
| feat = [x] | |
| for i in range(len(self.vgg_layers)): | |
| x = self.vgg_layers[i](x) | |
| if i in self.capture_layers: | |
| feat.append(x) | |
| return feat | |
| def forward(self, x): | |
| if self.space != 'vgg': | |
| x = (x + 1.) / 2. | |
| x = x - (torch.Tensor([0.485, 0.456, 0.406]).to(x.device).view(1, -1, 1, 1)) | |
| x = x / (torch.Tensor([0.229, 0.224, 0.225]).to(x.device).view(1, -1, 1, 1)) | |
| feat = self.forward_base(x) | |
| return feat | |
| def forward_samples_hypercolumn(self, X, samps=100): | |
| feat = self.forward(X) | |
| xx, xy = np.meshgrid(np.arange(X.shape[2]), np.arange(X.shape[3])) | |
| xx = np.expand_dims(xx.flatten(), 1) | |
| xy = np.expand_dims(xy.flatten(), 1) | |
| xc = np.concatenate([xx, xy], 1) | |
| samples = min(samps, xc.shape[0]) | |
| np.random.shuffle(xc) | |
| xx = xc[:samples, 0] | |
| yy = xc[:samples, 1] | |
| feat_samples = [] | |
| for i in range(len(feat)): | |
| layer_feat = feat[i] | |
| # hack to detect lower resolution | |
| if i > 0 and feat[i].size(2) < feat[i - 1].size(2): | |
| xx = xx / 2.0 | |
| yy = yy / 2.0 | |
| xx = np.clip(xx, 0, layer_feat.shape[2] - 1).astype(np.int32) | |
| yy = np.clip(yy, 0, layer_feat.shape[3] - 1).astype(np.int32) | |
| features = layer_feat[:, :, xx[range(samples)], yy[range(samples)]] | |
| feat_samples.append(features.clone().detach()) | |
| feat = torch.cat(feat_samples, 1) | |
| return feat | |
| class StyleLoss: | |
| def spatial_feature_extract(self, feat_result, feat_content, xx, xy): | |
| l2, l3 = [], [] | |
| device = feat_result[0].device | |
| # for each extracted layer | |
| for i in range(len(feat_result)): | |
| fr = feat_result[i] | |
| fc = feat_content[i] | |
| # hack to detect reduced scale | |
| if i > 0 and feat_result[i - 1].size(2) > feat_result[i].size(2): | |
| xx = xx / 2.0 | |
| xy = xy / 2.0 | |
| # go back to ints and get residual | |
| xxm = np.floor(xx).astype(np.float32) | |
| xxr = xx - xxm | |
| xym = np.floor(xy).astype(np.float32) | |
| xyr = xy - xym | |
| # do bilinear resample | |
| w00 = torch.from_numpy((1. - xxr) * (1. - xyr)).float().view(1, 1, -1, 1).to(device) | |
| w01 = torch.from_numpy((1. - xxr) * xyr).float().view(1, 1, -1, 1).to(device) | |
| w10 = torch.from_numpy(xxr * (1. - xyr)).float().view(1, 1, -1, 1).to(device) | |
| w11 = torch.from_numpy(xxr * xyr).float().view(1, 1, -1, 1).to(device) | |
| xxm = np.clip(xxm.astype(np.int32), 0, fr.size(2) - 1) | |
| xym = np.clip(xym.astype(np.int32), 0, fr.size(3) - 1) | |
| s00 = xxm * fr.size(3) + xym | |
| s01 = xxm * fr.size(3) + np.clip(xym + 1, 0, fr.size(3) - 1) | |
| s10 = np.clip(xxm + 1, 0, fr.size(2) - 1) * fr.size(3) + (xym) | |
| s11 = np.clip(xxm + 1, 0, fr.size(2) - 1) * fr.size(3) + np.clip(xym + 1, 0, fr.size(3) - 1) | |
| fr = fr.view(1, fr.size(1), fr.size(2) * fr.size(3), 1) | |
| fr = fr[:, :, s00, :].mul_(w00).add_(fr[:, :, s01, :].mul_(w01)).add_(fr[:, :, s10, :].mul_(w10)).add_( | |
| fr[:, :, s11, :].mul_(w11)) | |
| fc = fc.view(1, fc.size(1), fc.size(2) * fc.size(3), 1) | |
| fc = fc[:, :, s00, :].mul_(w00).add_(fc[:, :, s01, :].mul_(w01)).add_(fc[:, :, s10, :].mul_(w10)).add_( | |
| fc[:, :, s11, :].mul_(w11)) | |
| l2.append(fr) | |
| l3.append(fc) | |
| x_st = torch.cat([li.contiguous() for li in l2], 1) | |
| c_st = torch.cat([li.contiguous() for li in l3], 1) | |
| xx = torch.from_numpy(xx).view(1, 1, x_st.size(2), 1).float().to(device) | |
| yy = torch.from_numpy(xy).view(1, 1, x_st.size(2), 1).float().to(device) | |
| x_st = torch.cat([x_st, xx, yy], 1) | |
| c_st = torch.cat([c_st, xx, yy], 1) | |
| return x_st, c_st | |
| def rgb_to_yuv(self, rgb): | |
| C = torch.Tensor( | |
| [[0.577350, 0.577350, 0.577350], [-0.577350, 0.788675, -0.211325], [-0.577350, -0.211325, 0.788675]] | |
| ).to(rgb.device) | |
| yuv = torch.mm(C, rgb) | |
| return yuv | |
| def pairwise_distances_cos(self, x, y): | |
| x_norm = torch.sqrt((x ** 2).sum(1).view(-1, 1)) | |
| y_t = torch.transpose(y, 0, 1) | |
| y_norm = torch.sqrt((y ** 2).sum(1).view(1, -1)) | |
| dist = 1. - torch.mm(x, y_t) / x_norm / y_norm | |
| return dist | |
| def pairwise_distances_sq_l2(self, x, y): | |
| x_norm = (x ** 2).sum(1).view(-1, 1) | |
| y_t = torch.transpose(y, 0, 1) | |
| y_norm = (y ** 2).sum(1).view(1, -1) | |
| dist = x_norm + y_norm - 2.0 * torch.mm(x, y_t) | |
| return torch.clamp(dist, 1e-5, 1e5) / x.size(1) | |
| def distmat(self, x, y, cos_d=True): | |
| if cos_d: | |
| M = self.pairwise_distances_cos(x, y) | |
| else: | |
| M = torch.sqrt(self.pairwise_distances_sq_l2(x, y)) | |
| return M | |
| def style_loss(self, X, Y): | |
| d = X.shape[1] | |
| if d == 3: | |
| X = self.rgb_to_yuv(X.transpose(0, 1).contiguous().view(d, -1)).transpose(0, 1) | |
| Y = self.rgb_to_yuv(Y.transpose(0, 1).contiguous().view(d, -1)).transpose(0, 1) | |
| else: | |
| X = X.transpose(0, 1).contiguous().view(d, -1).transpose(0, 1) | |
| Y = Y.transpose(0, 1).contiguous().view(d, -1).transpose(0, 1) | |
| # Relaxed EMD | |
| CX_M = self.distmat(X, Y, cos_d=True) | |
| if d == 3: | |
| CX_M = CX_M + self.distmat(X, Y, cos_d=False) | |
| m1, m1_inds = CX_M.min(1) | |
| m2, m2_inds = CX_M.min(0) | |
| remd = torch.max(m1.mean(), m2.mean()) | |
| return remd | |
| def moment_loss(self, X, Y, moments=[1, 2]): | |
| loss = 0. | |
| X = X.squeeze().t() | |
| Y = Y.squeeze().t() | |
| mu_x = torch.mean(X, 0, keepdim=True) | |
| mu_y = torch.mean(Y, 0, keepdim=True) | |
| mu_d = torch.abs(mu_x - mu_y).mean() | |
| if 1 in moments: | |
| loss = loss + mu_d | |
| if 2 in moments: | |
| X_c = X - mu_x | |
| Y_c = Y - mu_y | |
| X_cov = torch.mm(X_c.t(), X_c) / (X.shape[0] - 1) | |
| Y_cov = torch.mm(Y_c.t(), Y_c) / (Y.shape[0] - 1) | |
| D_cov = torch.abs(X_cov - Y_cov).mean() | |
| loss = loss + D_cov | |
| return loss | |
| def forward(self, feat_result, feat_content, feat_style, indices, content_weight, moment_weight=1.0): | |
| # spatial feature extract | |
| num_locations = 1024 | |
| spatial_result, spatial_content = self.spatial_feature_extract( | |
| feat_result, feat_content, indices[0][:num_locations], indices[1][:num_locations] | |
| ) | |
| # loss_content = content_loss(spatial_result, spatial_content) | |
| d = feat_style.shape[1] | |
| spatial_style = feat_style.view(1, d, -1, 1) | |
| feat_max = 3 + 2 * 64 + 128 * 2 + 256 * 3 + 512 * 2 # (sum of all extracted channels) | |
| loss_remd = self.style_loss(spatial_result[:, :feat_max, :, :], spatial_style[:, :feat_max, :, :]) | |
| loss_moment = self.moment_loss(spatial_result[:, :-2, :, :], | |
| spatial_style, | |
| moments=[1, 2]) # -2 is so that it can fit? | |
| # palette matching | |
| content_weight_frac = 1. / max(content_weight, 1.) | |
| loss_moment += content_weight_frac * self.style_loss(spatial_result[:, :3, :, :], spatial_style[:, :3, :, :]) | |
| loss_style = loss_remd + moment_weight * loss_moment | |
| # print(f'Style: {loss_style.item():.3f}, Content: {loss_content.item():.3f}') | |
| style_weight = 1.0 + moment_weight | |
| loss_total = (loss_style) / (content_weight + style_weight) | |
| return loss_total | |
| def sample_indices(feat_content, feat_style): | |
| const = 128 ** 2 # 32k or so | |
| big_size = feat_content.shape[2] * feat_content.shape[3] # num feaxels | |
| stride_x = int(max(math.floor(math.sqrt(big_size // const)), 1)) | |
| offset_x = np.random.randint(stride_x) | |
| stride_y = int(max(math.ceil(math.sqrt(big_size // const)), 1)) | |
| offset_y = np.random.randint(stride_y) | |
| xx, xy = np.meshgrid( | |
| np.arange(feat_content.shape[2])[offset_x::stride_x], | |
| np.arange(feat_content.shape[3])[offset_y::stride_y] | |
| ) | |
| xx = xx.flatten() | |
| xy = xy.flatten() | |
| return xx, xy | |