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 |