Generater.py

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