Update Generater.py

Generater.py

@@ -1,42 +1,240 @@
-import cv2
+
 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
-from Generater import Generater
-import gradio as gr
-
-
-generator = Generater()
-# # oslist =os.listdir("纹理")
-# # print(oslist)
-G_path ='Gmodel_state33003.pdparams'
-layer_state_dictg = paddle.load(G_path)
-generator.set_state_dict(layer_state_dictg)#导入训练好的参数文件
-
-def style_transfer(content_img,style_img):
-    g_input = content_img.astype('float32') / 127.5 - 1  # 归一化
-    g_input = g_input[np.newaxis, ...].transpose(0, 3, 1, 2)  # NHWC -> NCHW
-    g_input = paddle.to_tensor(g_input)  # numpy -> tensor
-    h,w = g_input.shape[-2:]
-    p = max([h,w])
-    g_input = F.interpolate(g_input,scale_factor=(256/p))
-
-    g_input_s = style_img.astype('float32') / 127.5 - 1  # 归一化
-    g_input_s = g_input_s[np.newaxis, ...].transpose(0, 3, 1, 2)  # NHWC -> NCHW
-    g_input_s = paddle.to_tensor(g_input_s)  # numpy -> tensor
-    h,w = g_input_s.shape[-2:]
-    p = max([h,w])
-    g_input_s = F.interpolate(g_input_s,scale_factor=(256/p))
-
-    i = paddle.to_tensor([1.])
-    g_output = generator(g_input,g_input_s,i)
-    g_output = g_output.detach().numpy()                      # tensor -> numpy
-    g_output = g_output.transpose(0, 2, 3, 1)[0]             # NCHW -> NHWC
-    g_output = (g_output+1) *127.5                        # 反归一化
-    g_output = g_output.astype(np.uint8)
-    output = g_output
-    # cv2.imwrite(os.path.join("./test", str(i.numpy()[0])+'qt.png'), g_output)#保存图片到本地
-    return output
-
-interface = gr.Interface(fn=style_transfer, inputs=["image","image"], outputs="image")
-interface.launch(share=True)
+# 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