Update Generater.py

Generater.py

@@ -1,240 +1,42 @@
-
+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
-# 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
+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)