import numpy as np # import os import paddle import paddle.optimizer import paddle.nn as nn # from tqdm import tqdm # from paddle.io import Dataset # from paddle.io import DataLoader import paddle.nn.functional as F # import paddle.tensor as tensor class VGG19(nn.Layer): cfg = [ 64, 64, 'M', 128, 128, 'M', 256, 256, 256, 256, 'M', 512, 512, 512, 512,'M', 512, 512, 512, 512, 'M'] def __init__(self, output_index: int = 26) -> None: super().__init__() # arch = 'caffevgg19' # weights_path = get_path_from_url(model_urls[arch][0], # model_urls[arch][1]) data_dict: dict = np.load("./vgg19_no_fc.npy", encoding='latin1', allow_pickle=True).item() self.features = self.make_layers(self.cfg, data_dict) del data_dict self.features = nn.Sequential(*self.features.sublayers()[:output_index]) mean = paddle.to_tensor([103.939, 116.779, 123.68]) self.mean = mean.unsqueeze(0).unsqueeze(-1).unsqueeze(-1) def _process(self, x): rgb = (x * 0.5 + 0.5) * 255 # value to 255 bgr = paddle.stack((rgb[:, 2, :, :], rgb[:, 1, :, :], rgb[:, 0, :, :]), 1) # rgb to bgr return bgr - self.mean # vgg norm def _forward_impl(self, x): x = self._process(x) # NOTE get output with out relu activation x = self.features(x) return x def forward(self, x): return self._forward_impl(x) @staticmethod def get_conv_filter(data_dict, name): return data_dict[name][0] @staticmethod def get_bias(data_dict, name): return data_dict[name][1] @staticmethod def get_fc_weight(data_dict, name): return data_dict[name][0] def make_layers(self, cfg, data_dict, batch_norm=False) -> nn.Sequential: layers = [] in_channels = 3 block = 1 number = 1 for v in cfg: if v == 'M': layers += [nn.MaxPool2D(kernel_size=2, stride=2)] block += 1 number = 1 else: conv2d = nn.Conv2D(in_channels, v, kernel_size=3, padding=1) """ set value """ weight = paddle.to_tensor( self.get_conv_filter(data_dict, f'conv{block}_{number}')) weight = weight.transpose((3, 2, 0, 1)) bias = paddle.to_tensor( self.get_bias(data_dict, f'conv{block}_{number}')) conv2d.weight.set_value(weight) conv2d.bias.set_value(bias) number += 1 if batch_norm: layers += [conv2d, nn.BatchNorm2D(v), nn.ReLU()] else: layers += [conv2d, nn.ReLU()] in_channels = v # print("number",block) return nn.Sequential(*layers) class InvertedresBlock(nn.Layer): def __init__(self, in_channels: int, expansion: float, out_channels: int, bias_attr=False): super().__init__() self.in_channels = in_channels self.expansion = expansion self.out_channels = out_channels self.bottle_channels = round(self.expansion * self.out_channels) self.body = nn.Sequential( # pw Conv2DNormLReLU(self.in_channels, self.bottle_channels, kernel_size=1, bias_attr=bias_attr), # dw nn.Conv2D(self.bottle_channels, self.bottle_channels, kernel_size=3, stride=1, padding=0, groups=self.bottle_channels, bias_attr=True), nn.GroupNorm(1, self.bottle_channels), nn.LeakyReLU(0.2), # pw & linear nn.Conv2D(self.bottle_channels, self.out_channels, kernel_size=1, padding=0, bias_attr=False), nn.GroupNorm(1, self.out_channels), ) def forward(self, x0): x = self.body(x0) if self.in_channels == self.out_channels: out = paddle.add(x0, x) else: out = x return x class Conv2DNormLReLU(nn.Layer): def __init__(self, in_channels: int, out_channels: int, kernel_size: int = 3, stride: int = 1, padding: int = 1, bias_attr=False) -> None: super().__init__() self.conv = nn.Conv2D(in_channels, out_channels, kernel_size, stride, padding, bias_attr=bias_attr) # NOTE layer norm is crucial for animegan! self.norm = nn.GroupNorm(1, out_channels) self.lrelu = nn.LeakyReLU(0.2) def forward(self, x): x = self.conv(x) x = self.norm(x) x = self.lrelu(x) return x class Generater(nn.Layer): def __init__(self): super().__init__() self.VGG = VGG19() self.A = nn.Sequential(InvertedresBlock(512, 2, 256), InvertedresBlock(256, 2, 256), InvertedresBlock(256, 2, 256), InvertedresBlock(256, 2, 256), Conv2DNormLReLU(256, 128)) self.B = nn.Sequential(nn.Upsample(scale_factor=2, mode='bilinear'), Conv2DNormLReLU(128, 128), Conv2DNormLReLU(128, 128)) self.C = nn.Sequential(nn.Upsample(scale_factor=2, mode='bilinear'), Conv2DNormLReLU(128, 128), Conv2DNormLReLU(128, 128)) self.D = nn.Sequential(nn.Upsample(scale_factor=2, mode='bilinear'), Conv2DNormLReLU(128, 64), Conv2DNormLReLU(64, 64), Conv2DNormLReLU(64, 32, 7, padding=3)) self.out = nn.Sequential(nn.Conv2D(32, 3, 1, bias_attr=False), nn.Tanh()) # ,nn.Sigmoid()) def style_projection(self,content_feature,style_feature,alpha = 0.7): def scatter_numpy(dim, index, src): dst = src.copy() idx_xsection_shape = index.shape[:dim] + index.shape[dim + 1:] # print("idx_xsection_shape",idx_xsection_shape)#(b,c) dst_xsection_shape = dst.shape[:dim] + dst.shape[dim + 1:] def make_slice(arr, dim, i): slc = [slice(None)] * arr.ndim slc[dim] = i return tuple(slc) # We use index and dim parameters to create idx # idx is in a form that can be used as a NumPy advanced index for scattering of src param. idx = [[ *np.indices(idx_xsection_shape).reshape(index.ndim - 1, -1), index[make_slice(index, dim, i)].reshape(1, -1)[0] ] for i in range(index.shape[dim])] idx = list(np.concatenate(idx, axis=1)) # print("idx",idx) # idx.insert(dim, idx.pop()) if not np.isscalar(src): src_idx = list(idx)#使idx和src_idx并不是同一个内存空间 src_idx.pop(dim) src_idx.insert(dim, np.repeat(np.arange(index.shape[dim]), np.prod(idx_xsection_shape))) dst[tuple(idx)] = src[tuple(src_idx)] else: dst[idx] = src return dst b,c,h,w = content_feature.shape style_feature = F.interpolate(x=style_feature, size=content_feature.shape[-2:],mode="BILINEAR") content_feat = content_feature.reshape([b,c,h*w]).numpy() style_feat = style_feature.reshape([b,c,h*w]).numpy() # print("content_feat",content_feat.shape,b,c) # content_feat = np.reshape(content_feat, (b,c, -1))#(b,c,-1) # style_feat = np.reshape(style_feat, (b,c, -1))#(b,c,-1) # print(content_feat) content_feat_index = np.argsort(content_feat, axis=2) style_feat = np.sort(style_feat, axis=2) # print("content_feat_index",content_feat_index) # print("style_feat",style_feat) fr_feat = scatter_numpy(dim=2, index=content_feat_index, src=style_feat) fr_feat = fr_feat * alpha + content_feat * (1 - alpha) fr_feat = np.reshape(fr_feat, (b,c,h,w)) fr_feat = paddle.to_tensor(fr_feat) return fr_feat # @paddle.jit.to_static def forward(self,real_image,style_image,alpha): alpha = alpha.numpy()[0] # print("real_image",real_image.shape) content_feature = self.VGG(real_image) # print("content_feat",content_feature.shape) style_feature = self.VGG(style_image) fr_feat = self.style_projection(content_feature,style_feature,alpha) a = self.A(fr_feat) b = self.B(a) c = self.C(b) d = self.D(c) out = self.out(d) return out