Spaces:
Runtime error
Runtime error
| 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) | |
| def get_conv_filter(data_dict, name): | |
| return data_dict[name][0] | |
| def get_bias(data_dict, name): | |
| return data_dict[name][1] | |
| 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 |