init

.gitignore

@@ -0,0 +1,5 @@
+myvenv
+myvenv/**/*
+__pycache__
+flagged
+*.pth

README.md

@@ -1,12 +1 @@
----
-title: Pixelization
-emoji: 🚀
-colorFrom: blue
-colorTo: gray
-sdk: gradio
-sdk_version: 3.16.2
-app_file: app.py
-pinned: false
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+...

app.py

@@ -1,8 +1,76 @@
 import gradio as gr
+import functools
+from pixelization import Model
+import torch
+import argparse
+import huggingface_hub
+import os
 
-def greet(name):
-    return "Hello " + name + "!!"
+TOKEN = "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"
 
-iface = gr.Interface(fn=greet, inputs="text", outputs="text")
-iface.launch()
-					
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--theme', type=str, default='default')
+    parser.add_argument('--live', action='store_true')
+    parser.add_argument('--share', action='store_true')
+    parser.add_argument('--port', type=int)
+    parser.add_argument('--disable-queue',
+                        dest='enable_queue',
+                        action='store_false')
+    parser.add_argument('--allow-flagging', type=str, default='never')
+    return parser.parse_args()
+
+def main():
+  args = parse_args()
+
+
+  # DL MODEL
+  # PIX_MODEL
+  os.environ['PIX_MODEL'] = huggingface_hub.hf_hub_download("NoCrypt/pixelization_models", "pixelart_vgg19.pth", token=TOKEN);
+  # NET_MODEL
+  os.environ['NET_MODEL'] = huggingface_hub.hf_hub_download("NoCrypt/pixelization_models", "160_net_G_A.pth", token=TOKEN);
+  # ALIAS_MODEL
+  os.environ['ALIAS_MODEL'] = huggingface_hub.hf_hub_download("NoCrypt/pixelization_models", "alias_net.pth", token=TOKEN);
+  
+  # # For local testing
+  # # PIX_MODEL
+  # os.environ['PIX_MODEL'] = "pixelart_vgg19.pth"
+  # # NET_MODEL
+  # os.environ['NET_MODEL'] = "160_net_G_A.pth"
+  # # ALIAS_MODEL
+  # os.environ['ALIAS_MODEL'] = "alias_net.pth"
+
+
+  use_cpu = True
+  m = Model(device = "cpu" if use_cpu else "cuda")
+  m.load()
+
+  # To use GPU: Change use_cpu to false, and checkout my comment on networks.py at line 107 & 108
+  # + Use torch with cuda support (Change in requirements.txt)
+
+  gr.Interface(m.pixelize_modified,
+             [
+              gr.components.Image(type='pil', label='Input'),
+              gr.components.Slider(minimum=1, maximum=16, value=4, step=1, label='Pixel Size'),
+              gr.components.Checkbox(True, label="Upscale after")
+             ],
+              gr.components.Image(type='pil', label='Output'),
+              title="Pixelization",
+              description='''
+Demo for [WuZongWei6/Pixelization](https://github.com/WuZongWei6/Pixelization)
+
+Models that are used is private to comply with License.
+
+
+              ''',
+              theme=args.theme,
+              allow_flagging=args.allow_flagging,
+              live=args.live, 
+              ).launch(
+                enable_queue=args.enable_queue,
+                server_port=args.port,
+                share=args.share,
+              )
+
+if __name__ == '__main__':
+    main()

models/basic_layer.py

@@ -0,0 +1,429 @@
+import torch.nn as nn
+import torch
+import torch.nn.functional as F
+import numpy as np
+
+class ModulationConvBlock(nn.Module):
+    def __init__(self, input_dim, output_dim, kernel_size, stride=1,
+                 padding=0, norm='none', activation='relu', pad_type='zero'):
+        super(ModulationConvBlock, self).__init__()
+        self.in_c = input_dim
+        self.out_c = output_dim
+        self.ksize = kernel_size
+        self.stride = 1
+        self.padding = kernel_size // 2
+
+        self.eps = 1e-8
+        weight_shape = (output_dim, input_dim, kernel_size, kernel_size)
+        fan_in = kernel_size * kernel_size *input_dim
+        wscale = 1.0/np.sqrt(fan_in)
+
+        self.weight = nn.Parameter(torch.randn(*weight_shape))
+        self.wscale = wscale
+
+        self.bias = nn.Parameter(torch.zeros(output_dim))
+
+        self.activate = nn.LeakyReLU(negative_slope=0.2, inplace=True)
+        self.activate_scale = np.sqrt(2.0)
+
+    def forward(self, x, code):
+        batch,in_channel,height,width = x.shape
+        weight = self.weight * self.wscale
+        _weight = weight.view(1, self.ksize, self.ksize, self.in_c, self.out_c)
+        _weight = _weight * code.view(batch, 1, 1, self.in_c, 1)
+        # demodulation
+        _weight_norm = torch.sqrt(torch.sum(_weight ** 2, dim=[1, 2, 3]) + self.eps)
+        _weight = _weight / _weight_norm.view(batch, 1, 1, 1, self.out_c)
+        # fused_modulate
+        x = x.view(1, batch * self.in_c, x.shape[2], x.shape[3])
+        weight = _weight.permute(1, 2, 3, 0, 4).reshape(
+            self.ksize, self.ksize, self.in_c, batch * self.out_c)
+        # not use_conv2d_transpose
+        weight = weight.permute(3, 2, 0, 1)
+        x = F.conv2d(x,
+                     weight=weight,
+                     bias=None,
+                     stride=self.stride,
+                     padding=self.padding,
+                     groups=(batch if True else 1))
+
+        if True:#self.fused_modulate:
+            x = x.view(batch, self.out_c, height, width)
+        x = x+self.bias.view(1,-1,1,1)
+        x = self.activate(x)*self.activate_scale
+        return x
+
+
+class AliasConvBlock(nn.Module):
+    def __init__(self, input_dim, output_dim, kernel_size, stride,
+                 padding=0, norm='none', activation='relu', pad_type='zero'):
+        super(AliasConvBlock, self).__init__()
+        self.use_bias = True
+        # initialize padding
+        if pad_type == 'reflect':
+            self.pad = nn.ReflectionPad2d(padding)
+        elif pad_type == 'replicate':
+            self.pad = nn.ReplicationPad2d(padding)
+        elif pad_type == 'zero':
+            self.pad = nn.ZeroPad2d(padding)
+        else:
+            assert 0, "Unsupported padding type: {}".format(pad_type)
+
+        # initialize normalization
+        norm_dim = output_dim
+        if norm == 'bn':
+            self.norm = nn.BatchNorm2d(norm_dim)
+        elif norm == 'in':
+            # self.norm = nn.InstanceNorm2d(norm_dim, track_running_stats=True)
+            self.norm = nn.InstanceNorm2d(norm_dim)
+        elif norm == 'ln':
+            self.norm = LayerNorm(norm_dim)
+        elif norm == 'adain':
+            self.norm = AdaptiveInstanceNorm2d(norm_dim)
+        elif norm == 'none' or norm == 'sn':
+            self.norm = None
+        else:
+            assert 0, "Unsupported normalization: {}".format(norm)
+
+        # initialize activation
+        if activation == 'relu':
+            self.activation = nn.ReLU(inplace=True)
+        elif activation == 'lrelu':
+            self.activation = nn.LeakyReLU(0.2, inplace=True)
+        elif activation == 'prelu':
+            self.activation = nn.PReLU()
+        elif activation == 'selu':
+            self.activation = nn.SELU(inplace=True)
+        elif activation == 'tanh':
+            self.activation = nn.Tanh()
+        elif activation == 'none':
+            self.activation = None
+        else:
+            assert 0, "Unsupported activation: {}".format(activation)
+
+        # initialize convolution
+        if norm == 'sn':
+            self.conv = nn.Conv2d(input_dim, output_dim, kernel_size, stride, bias=self.use_bias)
+
+        else:
+            self.conv = nn.Conv2d(input_dim, output_dim, kernel_size, stride, bias=self.use_bias)
+
+    def forward(self, x):
+        x = self.conv(self.pad(x))
+        if self.norm:
+            x = self.norm(x)
+        if self.activation:
+            x = self.activation(x)
+        return x
+
+class AliasResBlocks(nn.Module):
+    def __init__(self, num_blocks, dim, norm='in', activation='relu', pad_type='zero'):
+        super(AliasResBlocks, self).__init__()
+        self.model = []
+        for i in range(num_blocks):
+            self.model += [AliasResBlock(dim, norm=norm, activation=activation, pad_type=pad_type)]
+        self.model = nn.Sequential(*self.model)
+
+    def forward(self, x):
+        return self.model(x)
+class AliasResBlock(nn.Module):
+    def __init__(self, dim, norm='in', activation='relu', pad_type='zero'):
+        super(AliasResBlock, self).__init__()
+
+        model = []
+        model += [AliasConvBlock(dim, dim, 3, 1, 1, norm=norm, activation=activation, pad_type=pad_type)]
+        model += [AliasConvBlock(dim, dim, 3, 1, 1, norm=norm, activation='none', pad_type=pad_type)]
+        self.model = nn.Sequential(*model)
+
+    def forward(self, x):
+        residual = x
+        out = self.model(x)
+        out += residual
+        return out
+##################################################################################
+# Sequential Models
+##################################################################################
+class ResBlocks(nn.Module):
+    def __init__(self, num_blocks, dim, norm='in', activation='relu', pad_type='zero'):
+        super(ResBlocks, self).__init__()
+        self.model = []
+        for i in range(num_blocks):
+            self.model += [ResBlock(dim, norm=norm, activation=activation, pad_type=pad_type)]
+        self.model = nn.Sequential(*self.model)
+
+    def forward(self, x):
+        return self.model(x)
+
+
+class MLP(nn.Module):
+    def __init__(self, input_dim, output_dim, dim, n_blk, norm='none', activ='relu'):
+        super(MLP, self).__init__()
+        self.model = []
+        self.model += [linearBlock(input_dim, input_dim, norm=norm, activation=activ)]
+        self.model += [linearBlock(input_dim, dim, norm=norm, activation=activ)]
+        for i in range(n_blk - 2):
+            self.model += [linearBlock(dim, dim, norm=norm, activation=activ)]
+        self.model += [linearBlock(dim, output_dim, norm='none', activation='none')]  # no output activations
+        self.model = nn.Sequential(*self.model)
+
+    # def forward(self, style0, style1, a=0):
+    #     return self.model[3]((1 - a) * self.model[0:3](style0.view(style0.size(0), -1)) + a * self.model[0:3](
+    #         style1.view(style1.size(0), -1)))
+    def forward(self, style0, style1=None, a=0):
+        style1 = style0
+        return self.model[3]((1 - a) * self.model[0:3](style0.view(style0.size(0), -1)) + a * self.model[0:3](
+            style1.view(style1.size(0), -1)))
+##################################################################################
+# Basic Blocks
+##################################################################################
+class ResBlock(nn.Module):
+    def __init__(self, dim, norm='in', activation='relu', pad_type='zero'):
+        super(ResBlock, self).__init__()
+
+        model = []
+        model += [ConvBlock(dim, dim, 3, 1, 1, norm=norm, activation=activation, pad_type=pad_type)]
+        model += [ConvBlock(dim, dim, 3, 1, 1, norm=norm, activation='none', pad_type=pad_type)]
+        self.model = nn.Sequential(*model)
+
+    def forward(self, x):
+        residual = x
+        out = self.model(x)
+        out += residual
+        return out
+
+
+class ConvBlock(nn.Module):
+    def __init__(self, input_dim, output_dim, kernel_size, stride,
+                 padding=0, norm='none', activation='relu', pad_type='zero'):
+        super(ConvBlock, self).__init__()
+        self.use_bias = True
+        # initialize padding
+        if pad_type == 'reflect':
+            self.pad = nn.ReflectionPad2d(padding)
+        elif pad_type == 'replicate':
+            self.pad = nn.ReplicationPad2d(padding)
+        elif pad_type == 'zero':
+            self.pad = nn.ZeroPad2d(padding)
+        else:
+            assert 0, "Unsupported padding type: {}".format(pad_type)
+
+        # initialize normalization
+        norm_dim = output_dim
+        if norm == 'bn':
+            self.norm = nn.BatchNorm2d(norm_dim)
+        elif norm == 'in':
+            # self.norm = nn.InstanceNorm2d(norm_dim, track_running_stats=True)
+            self.norm = nn.InstanceNorm2d(norm_dim)
+        elif norm == 'ln':
+            self.norm = LayerNorm(norm_dim)
+        elif norm == 'adain':
+            self.norm = AdaptiveInstanceNorm2d(norm_dim)
+        elif norm == 'none' or norm == 'sn':
+            self.norm = None
+        else:
+            assert 0, "Unsupported normalization: {}".format(norm)
+
+        # initialize activation
+        if activation == 'relu':
+            self.activation = nn.ReLU(inplace=True)
+        elif activation == 'lrelu':
+            self.activation = nn.LeakyReLU(0.2, inplace=True)
+        elif activation == 'prelu':
+            self.activation = nn.PReLU()
+        elif activation == 'selu':
+            self.activation = nn.SELU(inplace=True)
+        elif activation == 'tanh':
+            self.activation = nn.Tanh()
+        elif activation == 'none':
+            self.activation = None
+        else:
+            assert 0, "Unsupported activation: {}".format(activation)
+
+        # initialize convolution
+        if norm == 'sn':
+            self.conv = nn.Conv2d(input_dim, output_dim, kernel_size, stride, bias=self.use_bias)
+
+        else:
+            self.conv = nn.Conv2d(input_dim, output_dim, kernel_size, stride, bias=self.use_bias)
+
+    def forward(self, x):
+        x = self.conv(self.pad(x))
+        if self.norm:
+            x = self.norm(x)
+        if self.activation:
+            x = self.activation(x)
+        return x
+
+class linearBlock(nn.Module):
+    def __init__(self, input_dim, output_dim, norm='none', activation='relu'):
+        super(linearBlock, self).__init__()
+        use_bias = True
+        # initialize fully connected layer
+        if norm == 'sn':
+            self.fc = SpectralNorm(nn.Linear(input_dim, output_dim, bias=use_bias))
+        else:
+            self.fc = nn.Linear(input_dim, output_dim, bias=use_bias)
+
+        # initialize normalization
+        norm_dim = output_dim
+        if norm == 'bn':
+            self.norm = nn.BatchNorm1d(norm_dim)
+        elif norm == 'in':
+            self.norm = nn.InstanceNorm1d(norm_dim)
+        elif norm == 'ln':
+            self.norm = LayerNorm(norm_dim)
+        elif norm == 'none' or norm == 'sn':
+            self.norm = None
+        else:
+            assert 0, "Unsupported normalization: {}".format(norm)
+
+        # initialize activation
+        if activation == 'relu':
+            self.activation = nn.ReLU(inplace=True)
+        elif activation == 'lrelu':
+            self.activation = nn.LeakyReLU(0.2, inplace=True)
+        elif activation == 'prelu':
+            self.activation = nn.PReLU()
+        elif activation == 'selu':
+            self.activation = nn.SELU(inplace=True)
+        elif activation == 'tanh':
+            self.activation = nn.Tanh()
+        elif activation == 'none':
+            self.activation = None
+        else:
+            assert 0, "Unsupported activation: {}".format(activation)
+
+    def forward(self, x):
+        out = self.fc(x)
+        if self.norm:
+            out = self.norm(out)
+        if self.activation:
+            out = self.activation(out)
+        return out
+##################################################################################
+# Normalization layers
+##################################################################################
+class AdaptiveInstanceNorm2d(nn.Module):
+    def __init__(self, num_features, eps=1e-5, momentum=0.1):
+        super(AdaptiveInstanceNorm2d, self).__init__()
+        self.num_features = num_features
+        self.eps = eps
+        self.momentum = momentum
+        # weight and bias are dynamically assigned
+        self.weight = None
+        self.bias = None
+        # just dummy buffers, not used
+        self.register_buffer('running_mean', torch.zeros(num_features))
+        self.register_buffer('running_var', torch.ones(num_features))
+
+    def forward(self, x):
+        assert self.weight is not None and self.bias is not None, "Please assign weight and bias before calling AdaIN!"
+        b, c = x.size(0), x.size(1)
+        running_mean = self.running_mean.repeat(b)
+        running_var = self.running_var.repeat(b)
+
+        # Apply instance norm
+        x_reshaped = x.contiguous().view(1, b * c, *x.size()[2:])
+
+        out = F.batch_norm(
+            x_reshaped, running_mean, running_var, self.weight, self.bias,
+            True, self.momentum, self.eps)
+
+        return out.view(b, c, *x.size()[2:])
+
+    def __repr__(self):
+        return self.__class__.__name__ + '(' + str(self.num_features) + ')'
+
+
+class LayerNorm(nn.Module):
+    def __init__(self, num_features, eps=1e-5, affine=True):
+        super(LayerNorm, self).__init__()
+        self.num_features = num_features
+        self.affine = affine
+        self.eps = eps
+
+        if self.affine:
+            self.gamma = nn.Parameter(torch.Tensor(num_features).uniform_())
+            self.beta = nn.Parameter(torch.zeros(num_features))
+
+    def forward(self, x):
+        shape = [-1] + [1] * (x.dim() - 1)
+        # print(x.size())
+        if x.size(0) == 1:
+            # These two lines run much faster in pytorch 0.4 than the two lines listed below.
+            mean = x.view(-1).mean().view(*shape)
+            std = x.view(-1).std().view(*shape)
+        else:
+            mean = x.view(x.size(0), -1).mean(1).view(*shape)
+            std = x.view(x.size(0), -1).std(1).view(*shape)
+
+        x = (x - mean) / (std + self.eps)
+
+        if self.affine:
+            shape = [1, -1] + [1] * (x.dim() - 2)
+            x = x * self.gamma.view(*shape) + self.beta.view(*shape)
+        return x
+
+
+def l2normalize(v, eps=1e-12):
+    return v / (v.norm() + eps)
+
+
+class SpectralNorm(nn.Module):
+    """
+    Based on the paper "Spectral Normalization for Generative Adversarial Networks" by Takeru Miyato, Toshiki Kataoka, Masanori Koyama, Yuichi Yoshida
+    and the Pytorch implementation https://github.com/christiancosgrove/pytorch-spectral-normalization-gan
+    """
+
+    def __init__(self, module, name='weight', power_iterations=1):
+        super(SpectralNorm, self).__init__()
+        self.module = module
+        self.name = name
+        self.power_iterations = power_iterations
+        if not self._made_params():
+            self._make_params()
+
+    def _update_u_v(self):
+        u = getattr(self.module, self.name + "_u")
+        v = getattr(self.module, self.name + "_v")
+        w = getattr(self.module, self.name + "_bar")
+
+        height = w.data.shape[0]
+        for _ in range(self.power_iterations):
+            v.data = l2normalize(torch.mv(torch.t(w.view(height, -1).data), u.data))
+            u.data = l2normalize(torch.mv(w.view(height, -1).data, v.data))
+
+        # sigma = torch.dot(u.data, torch.mv(w.view(height,-1).data, v.data))
+        sigma = u.dot(w.view(height, -1).mv(v))
+        setattr(self.module, self.name, w / sigma.expand_as(w))
+
+    def _made_params(self):
+        try:
+            u = getattr(self.module, self.name + "_u")
+            v = getattr(self.module, self.name + "_v")
+            w = getattr(self.module, self.name + "_bar")
+            return True
+        except AttributeError:
+            return False
+
+    def _make_params(self):
+        w = getattr(self.module, self.name)
+
+        height = w.data.shape[0]
+        width = w.view(height, -1).data.shape[1]
+
+        u = nn.Parameter(w.data.new(height).normal_(0, 1), requires_grad=False)
+        v = nn.Parameter(w.data.new(width).normal_(0, 1), requires_grad=False)
+        u.data = l2normalize(u.data)
+        v.data = l2normalize(v.data)
+        w_bar = nn.Parameter(w.data)
+
+        del self.module._parameters[self.name]
+
+        self.module.register_parameter(self.name + "_u", u)
+        self.module.register_parameter(self.name + "_v", v)
+        self.module.register_parameter(self.name + "_bar", w_bar)
+
+    def forward(self, *args):
+        self._update_u_v()
+        return self.module.forward(*args)

models/c2pDis.py

@@ -0,0 +1,313 @@
+from .basic_layer import *
+import math
+from torch.nn import Parameter
+#from pytorch_metric_learning import losses
+
+'''
+Margin code is borrowed from https://github.com/MuggleWang/CosFace_pytorch and https://github.com/wujiyang/Face_Pytorch.
+'''
+def cosine_sim(x1, x2, dim=1, eps=1e-8):
+    ip = torch.mm(x1, x2.t())  #  w  7*512
+    w1 = torch.norm(x1, 2, dim)
+    w2 = torch.norm(x2, 2, dim)
+    return ip / torch.ger(w1,w2).clamp(min=eps)
+
+class MarginCosineProduct(nn.Module):
+    r"""Implement of large margin cosine distance: :
+    Args:
+        in_features: size of each input sample
+        out_features: size of each output sample
+        s: norm of input feature
+        m: margin
+    """
+
+    def __init__(self, in_features, out_features, s=30.0, m=0.40):
+        super(MarginCosineProduct, self).__init__()
+        self.in_features = in_features
+        self.out_features = out_features
+        self.s = s
+        self.m = m
+        self.weight = Parameter(torch.Tensor(out_features, in_features)) # 7 512
+        nn.init.xavier_uniform_(self.weight)
+        #stdv = 1. / math.sqrt(self.weight.size(1))
+        #self.weight.data.uniform_(-stdv, stdv)
+
+    def forward(self, input, label):
+        cosine = cosine_sim(input, self.weight)  # 1*512  7*512
+        # cosine = F.linear(F.normalize(input), F.normalize(self.weight))
+        # --------------------------- convert label to one-hot ---------------------------
+        # https://discuss.pytorch.org/t/convert-int-into-one-hot-format/507
+        one_hot = torch.zeros_like(cosine)
+        one_hot.scatter_(1, label.view(-1, 1), 1.0)
+        # -------------torch.where(out_i = {x_i if condition_i else y_i) -------------
+        output = self.s * (cosine - one_hot * self.m)
+
+        return output
+
+    def __repr__(self):
+        return self.__class__.__name__ + '(' \
+               + 'in_features=' + str(self.in_features) \
+               + ', out_features=' + str(self.out_features) \
+               + ', s=' + str(self.s) \
+               + ', m=' + str(self.m) + ')'
+
+class ArcMarginProduct(nn.Module):
+    def __init__(self, in_feature=128, out_feature=10575, s=32.0, m=0.50, easy_margin=False):
+        super(ArcMarginProduct, self).__init__()
+        self.in_feature = in_feature
+        self.out_feature = out_feature
+        self.s = s
+        self.m = m
+        self.weight = Parameter(torch.Tensor(out_feature, in_feature))
+        nn.init.xavier_uniform_(self.weight)
+
+        self.easy_margin = easy_margin
+        self.cos_m = math.cos(m)
+        self.sin_m = math.sin(m)
+
+        # make the function cos(theta+m) monotonic decreasing while theta in [0°,180°]
+        self.th = math.cos(math.pi - m)
+        self.mm = math.sin(math.pi - m) * m
+
+    def forward(self, x, label):
+        # cos(theta)
+        cosine = F.linear(F.normalize(x), F.normalize(self.weight))
+        # cos(theta + m)
+        sine = torch.sqrt(1.0 - torch.pow(cosine, 2))
+        phi = cosine * self.cos_m - sine * self.sin_m
+
+        if self.easy_margin:
+            phi = torch.where(cosine > 0, phi, cosine)
+        else:
+            phi = torch.where((cosine - self.th) > 0, phi, cosine - self.mm)
+
+        #one_hot = torch.zeros(cosine.size(), device='cuda' if torch.cuda.is_available() else 'cpu')
+        one_hot = torch.zeros_like(cosine)
+        one_hot.scatter_(1, label.view(-1, 1), 1)
+        output = (one_hot * phi) + ((1.0 - one_hot) * cosine)
+        output = output * self.s
+
+        return output
+
+
+class MultiMarginProduct(nn.Module):
+    def __init__(self, in_feature=128, out_feature=10575, s=32.0, m1=0.20, m2=0.35, easy_margin=False):
+        super(MultiMarginProduct, self).__init__()
+        self.in_feature = in_feature
+        self.out_feature = out_feature
+        self.s = s
+        self.m1 = m1
+        self.m2 = m2
+        self.weight = Parameter(torch.Tensor(out_feature, in_feature))
+        nn.init.xavier_uniform_(self.weight)
+
+        self.easy_margin = easy_margin
+        self.cos_m1 = math.cos(m1)
+        self.sin_m1 = math.sin(m1)
+
+        # make the function cos(theta+m) monotonic decreasing while theta in [0°,180°]
+        self.th = math.cos(math.pi - m1)
+        self.mm = math.sin(math.pi - m1) * m1
+
+    def forward(self, x, label):
+        # cos(theta)
+        cosine = F.linear(F.normalize(x), F.normalize(self.weight))
+        # cos(theta + m1)
+        sine = torch.sqrt(1.0 - torch.pow(cosine, 2))
+        phi = cosine * self.cos_m1 - sine * self.sin_m1
+
+        if self.easy_margin:
+            phi = torch.where(cosine > 0, phi, cosine)
+        else:
+            phi = torch.where((cosine - self.th) > 0, phi, cosine - self.mm)
+
+
+        one_hot = torch.zeros_like(cosine)
+        one_hot.scatter_(1, label.view(-1, 1), 1)
+        output = (one_hot * phi) + ((1.0 - one_hot) * cosine) # additive angular margin
+        output = output - one_hot * self.m2 # additive cosine margin
+        output = output * self.s
+
+        return output
+
+
+class CPDis(nn.Module):
+    """PatchGAN."""
+    def __init__(self, image_size=256, conv_dim=64, repeat_num=3, norm='SN'):
+        super(CPDis, self).__init__()
+
+        layers = []
+        if norm == 'SN':
+            layers.append(spectral_norm(nn.Conv2d(3, conv_dim, kernel_size=4, stride=2, padding=1)))
+        else:
+            layers.append(nn.Conv2d(3, conv_dim, kernel_size=4, stride=2, padding=1))
+        layers.append(nn.LeakyReLU(0.01, inplace=True))
+
+        curr_dim = conv_dim
+        for i in range(1, repeat_num):
+            if norm == 'SN':
+                layers.append(spectral_norm(nn.Conv2d(curr_dim, curr_dim * 2, kernel_size=4, stride=2, padding=1)))
+            else:
+                layers.append(nn.Conv2d(curr_dim, curr_dim * 2, kernel_size=4, stride=2, padding=1))
+            layers.append(nn.LeakyReLU(0.01, inplace=True))
+            curr_dim = curr_dim * 2
+
+        # k_size = int(image_size / np.power(2, repeat_num))
+        if norm == 'SN':
+            layers.append(spectral_norm(nn.Conv2d(curr_dim, curr_dim * 2, kernel_size=4, stride=1, padding=1)))
+        else:
+            layers.append(nn.Conv2d(curr_dim, curr_dim * 2, kernel_size=4, stride=1, padding=1))
+        layers.append(nn.LeakyReLU(0.01, inplace=True))
+        curr_dim = curr_dim * 2
+
+        self.main = nn.Sequential(*layers)
+        if norm == 'SN':
+            self.conv1 = spectral_norm(nn.Conv2d(curr_dim, 1, kernel_size=4, stride=1, padding=1, bias=False))
+        else:
+            self.conv1 = nn.Conv2d(curr_dim, 1, kernel_size=4, stride=1, padding=1, bias=False)
+
+    def forward(self, x):
+        if x.ndim == 5:
+            x = x.squeeze(0)
+        assert x.ndim == 4, x.ndim
+        h = self.main(x)
+        # out_real = self.conv1(h)
+        out_makeup = self.conv1(h)
+        # return out_real.squeeze(), out_makeup.squeeze()
+        return out_makeup
+
+
+class CPDis_cls(nn.Module):
+    """PatchGAN."""
+    def __init__(self, image_size=256, conv_dim=64, repeat_num=3, norm='SN'):
+        super(CPDis_cls, self).__init__()
+
+        layers = []
+        if norm == 'SN':
+            layers.append(spectral_norm(nn.Conv2d(3, conv_dim, kernel_size=4, stride=2, padding=1)))
+        else:
+            layers.append(nn.Conv2d(3, conv_dim, kernel_size=4, stride=2, padding=1))
+        layers.append(nn.LeakyReLU(0.01, inplace=True))
+
+        curr_dim = conv_dim
+        for i in range(1, repeat_num):
+            if norm == 'SN':
+                layers.append(spectral_norm(nn.Conv2d(curr_dim, curr_dim * 2, kernel_size=4, stride=2, padding=1)))
+            else:
+                layers.append(nn.Conv2d(curr_dim, curr_dim * 2, kernel_size=4, stride=2, padding=1))
+            layers.append(nn.LeakyReLU(0.01, inplace=True))
+            curr_dim = curr_dim * 2
+
+        # k_size = int(image_size / np.power(2, repeat_num))
+        if norm == 'SN':
+            layers.append(spectral_norm(nn.Conv2d(curr_dim, curr_dim * 2, kernel_size=4, stride=1, padding=1)))
+        else:
+            layers.append(nn.Conv2d(curr_dim, curr_dim * 2, kernel_size=4, stride=1, padding=1))
+        layers.append(nn.LeakyReLU(0.01, inplace=True))
+        curr_dim = curr_dim * 2
+
+        self.main = nn.Sequential(*layers)
+        if norm == 'SN':
+            self.conv1 = spectral_norm(nn.Conv2d(curr_dim, 1, kernel_size=4, stride=1, padding=1, bias=False))
+            self.classifier_pool = nn.AdaptiveAvgPool2d(1)
+            self.classifier_conv = nn.Conv2d(512, 512, 1, 1, 0)
+            self.classifier = MarginCosineProduct(512,7)#ArcMarginProduct(512, 7)
+            print("Using Large Margin Cosine Loss.")
+
+        else:
+            self.conv1 = nn.Conv2d(curr_dim, 1, kernel_size=4, stride=1, padding=1, bias=False)
+
+    def forward(self, x, label):
+        if x.ndim == 5:
+            x = x.squeeze(0)
+        assert x.ndim == 4, x.ndim
+        h = self.main(x)  # ([1, 512, 31, 31])
+        #print(out_cls.shape)
+        out_cls = self.classifier_pool(h)
+        #print(out_cls.shape)
+        out_cls = self.classifier_conv(out_cls)
+        #print(out_cls.shape)
+        out_cls = torch.squeeze(out_cls, -1)
+        out_cls = torch.squeeze(out_cls, -1)
+        out_cls = self.classifier(out_cls, label)
+        out_makeup = self.conv1(h)  # torch.Size([1, 1, 30, 30])
+        # return out_real.squeeze(), out_makeup.squeeze()
+        return out_makeup, out_cls
+
+class SpectralNorm(object):
+    def __init__(self):
+        self.name = "weight"
+        # print(self.name)
+        self.power_iterations = 1
+
+    def compute_weight(self, module):
+        u = getattr(module, self.name + "_u")
+        v = getattr(module, self.name + "_v")
+        w = getattr(module, self.name + "_bar")
+
+        height = w.data.shape[0]
+        for _ in range(self.power_iterations):
+            v.data = l2normalize(torch.mv(torch.t(w.view(height, -1).data), u.data))
+            u.data = l2normalize(torch.mv(w.view(height, -1).data, v.data))
+        # sigma = torch.dot(u.data, torch.mv(w.view(height,-1).data, v.data))
+        sigma = u.dot(w.view(height, -1).mv(v))
+        return w / sigma.expand_as(w)
+
+    @staticmethod
+    def apply(module):
+        name = "weight"
+        fn = SpectralNorm()
+
+        try:
+            u = getattr(module, name + "_u")
+            v = getattr(module, name + "_v")
+            w = getattr(module, name + "_bar")
+        except AttributeError:
+            w = getattr(module, name)
+            height = w.data.shape[0]
+            width = w.view(height, -1).data.shape[1]
+            u = Parameter(w.data.new(height).normal_(0, 1), requires_grad=False)
+            v = Parameter(w.data.new(width).normal_(0, 1), requires_grad=False)
+            w_bar = Parameter(w.data)
+
+            # del module._parameters[name]
+
+            module.register_parameter(name + "_u", u)
+            module.register_parameter(name + "_v", v)
+            module.register_parameter(name + "_bar", w_bar)
+
+        # remove w from parameter list
+        del module._parameters[name]
+
+        setattr(module, name, fn.compute_weight(module))
+
+        # recompute weight before every forward()
+        module.register_forward_pre_hook(fn)
+
+        return fn
+
+    def remove(self, module):
+        weight = self.compute_weight(module)
+        delattr(module, self.name)
+        del module._parameters[self.name + '_u']
+        del module._parameters[self.name + '_v']
+        del module._parameters[self.name + '_bar']
+        module.register_parameter(self.name, Parameter(weight.data))
+
+    def __call__(self, module, inputs):
+        setattr(module, self.name, self.compute_weight(module))
+
+def spectral_norm(module):
+    SpectralNorm.apply(module)
+    return module
+
+def remove_spectral_norm(module):
+    name = 'weight'
+    for k, hook in module._forward_pre_hooks.items():
+        if isinstance(hook, SpectralNorm) and hook.name == name:
+            hook.remove(module)
+            del module._forward_pre_hooks[k]
+            return module
+
+    raise ValueError("spectral_norm of '{}' not found in {}"
+                     .format(name, module))

models/c2pGen.py

@@ -0,0 +1,266 @@
+from .basic_layer import *
+import torchvision.models as models
+import os
+
+
+
+class AliasNet(nn.Module):
+    def __init__(self, input_dim, output_dim, dim, n_downsample, n_res, activ='relu', pad_type='reflect'):
+        super(AliasNet, self).__init__()
+        self.RGBEnc = AliasRGBEncoder(input_dim, dim, n_downsample, n_res, "in", activ, pad_type=pad_type)
+        self.RGBDec = AliasRGBDecoder(self.RGBEnc.output_dim, output_dim, n_downsample, n_res, res_norm='in',
+                                      activ=activ, pad_type=pad_type)
+
+    def forward(self, x):
+        x = self.RGBEnc(x)
+        x = self.RGBDec(x)
+        return x
+
+
+class AliasRGBEncoder(nn.Module):
+    def __init__(self, input_dim, dim, n_downsample, n_res, norm, activ, pad_type):
+        super(AliasRGBEncoder, self).__init__()
+        self.model = []
+        self.model += [AliasConvBlock(input_dim, dim, 7, 1, 3, norm=norm, activation=activ, pad_type=pad_type)]
+        # downsampling blocks
+        for i in range(n_downsample):
+            self.model += [AliasConvBlock(dim, 2 * dim, 4, 2, 1, norm=norm, activation=activ, pad_type=pad_type)]
+            dim *= 2
+        # residual blocks
+        self.model += [AliasResBlocks(n_res, dim, norm=norm, activation=activ, pad_type=pad_type)]
+        self.model = nn.Sequential(*self.model)
+        self.output_dim = dim
+
+    def forward(self, x):
+        return self.model(x)
+
+
+class AliasRGBDecoder(nn.Module):
+    def __init__(self, dim, output_dim, n_upsample, n_res, res_norm, activ='relu', pad_type='zero'):
+        super(AliasRGBDecoder, self).__init__()
+        # self.model = []
+        # # AdaIN residual blocks
+        # self.model += [ResBlocks(n_res, dim, res_norm, activ, pad_type=pad_type)]
+        # # upsampling blocks
+        # for i in range(n_upsample):
+        #     self.model += [nn.Upsample(scale_factor=2, mode='nearest'),
+        #                    ConvBlock(dim, dim // 2, 5, 1, 2, norm='ln', activation=activ, pad_type=pad_type)]
+        #     dim //= 2
+        # # use reflection padding in the last conv layer
+        # self.model += [ConvBlock(dim, output_dim, 7, 1, 3, norm='none', activation='tanh', pad_type=pad_type)]
+        # self.model = nn.Sequential(*self.model)
+        self.Res_Blocks = AliasResBlocks(n_res, dim, res_norm, activ, pad_type=pad_type)
+        self.upsample_block1 = nn.Upsample(scale_factor=2, mode='nearest')
+        self.conv_1 = AliasConvBlock(dim, dim // 2, 5, 1, 2, norm='ln', activation=activ, pad_type=pad_type)
+        dim //= 2
+        self.upsample_block2 = nn.Upsample(scale_factor=2, mode='nearest')
+        self.conv_2 = AliasConvBlock(dim, dim // 2, 5, 1, 2, norm='ln', activation=activ, pad_type=pad_type)
+        dim //= 2
+        self.conv_3 = AliasConvBlock(dim, output_dim, 7, 1, 3, norm='none', activation='tanh', pad_type=pad_type)
+
+    def forward(self, x):
+        x = self.Res_Blocks(x)
+        # print(x.shape)
+        x = self.upsample_block1(x)
+        # print(x.shape)
+        x = self.conv_1(x)
+        # print(x_small.shape)
+        x = self.upsample_block2(x)
+        # print(x.shape)
+        x = self.conv_2(x)
+        # print(x_middle.shape)
+        x = self.conv_3(x)
+        # print(x_big.shape)
+        return x
+
+
+class C2PGen(nn.Module):
+    def __init__(self, input_dim, output_dim, dim, n_downsample, n_res, style_dim, mlp_dim, activ='relu', pad_type='reflect'):
+        super(C2PGen, self).__init__()
+        self.PBEnc = PixelBlockEncoder(input_dim, dim, style_dim, norm='none', activ=activ, pad_type=pad_type)
+        self.RGBEnc = RGBEncoder(input_dim, dim, n_downsample, n_res, "in", activ, pad_type=pad_type)
+        self.RGBDec = RGBDecoder(self.RGBEnc.output_dim, output_dim, n_downsample, n_res, res_norm='adain',
+                                      activ=activ, pad_type=pad_type)
+        self.MLP = MLP(style_dim, 2048, mlp_dim, 3, norm='none', activ=activ)
+
+    def forward(self, clipart, pixelart, s=1):
+        feature = self.RGBEnc(clipart)
+        code = self.PBEnc(pixelart)
+        result, cellcode = self.fuse(feature, code, s)
+        return result#, cellcode   #return cellcode when visualizing the cell size code
+
+    def fuse(self, content, style_code, s=1):
+        #print("MLP input:code's shape:", style_code.shape)
+        adain_params = self.MLP(style_code) * s # [batch,2048]
+        #print("MLP output:adain_params's shape", adain_params.shape)
+        #self.assign_adain_params(adain_params, self.RGBDec)
+        images = self.RGBDec(content, adain_params)
+        return images, adain_params
+
+    def assign_adain_params(self, adain_params, model):
+        # assign the adain_params to the AdaIN layers in model
+        for m in model.modules():
+            if m.__class__.__name__ == "AdaptiveInstanceNorm2d":
+                mean = adain_params[:, :m.num_features]
+                std = adain_params[:, m.num_features:2 * m.num_features]
+                m.bias = mean.contiguous().view(-1)
+                m.weight = std.contiguous().view(-1)
+                if adain_params.size(1) > 2 * m.num_features:
+                    adain_params = adain_params[:, 2 * m.num_features:]
+
+    def get_num_adain_params(self, model):
+        # return the number of AdaIN parameters needed by the model
+        num_adain_params = 0
+        for m in model.modules():
+            if m.__class__.__name__ == "AdaptiveInstanceNorm2d":
+                num_adain_params += 2 * m.num_features
+        return num_adain_params
+
+
+class PixelBlockEncoder(nn.Module):
+    def __init__(self, input_dim, dim, style_dim, norm, activ, pad_type):
+        super(PixelBlockEncoder, self).__init__()
+        vgg19 = models.vgg.vgg19()
+        vgg19.classifier._modules['6'] = nn.Linear(4096, 7, bias=True)
+        vgg19.load_state_dict(torch.load('./pixelart_vgg19.pth' if not os.environ['PIX_MODEL'] else os.environ['PIX_MODEL'], map_location=torch.device('cpu')))
+        self.vgg = vgg19.features
+        for p in self.vgg.parameters():
+            p.requires_grad = False
+        # vgg19 = models.vgg.vgg19(pretrained=False)
+        # vgg19.load_state_dict(torch.load('./vgg.pth'))
+        # self.vgg = vgg19.features
+        # for p in self.vgg.parameters():
+        #     p.requires_grad = False
+
+
+        self.conv1 = ConvBlock(input_dim, dim, 7, 1, 3, norm=norm, activation=activ, pad_type=pad_type)  # 3->64,concat
+        dim = dim * 2
+        self.conv2 = ConvBlock(dim, dim, 4, 2, 1, norm=norm, activation=activ, pad_type=pad_type)  # 128->128
+        dim = dim * 2
+        self.conv3 = ConvBlock(dim, dim, 4, 2, 1, norm=norm, activation=activ, pad_type=pad_type)  # 256->256
+        dim = dim * 2
+        self.conv4 = ConvBlock(dim, dim, 4, 2, 1, norm=norm, activation=activ, pad_type=pad_type)  # 512->512
+        dim = dim * 2
+
+        self.model = []
+        self.model += [nn.AdaptiveAvgPool2d(1)]  # global average pooling
+        self.model += [nn.Conv2d(dim, style_dim, 1, 1, 0)]
+        self.model = nn.Sequential(*self.model)
+        self.output_dim = dim
+
+    def get_features(self, image, model, layers=None):
+        if layers is None:
+            layers = {'0': 'conv1_1', '5': 'conv2_1', '10': 'conv3_1', '19': 'conv4_1'}
+        features = {}
+        x = image
+        # model._modules is a dictionary holding each module in the model
+        for name, layer in model._modules.items():
+            x = layer(x)
+            if name in layers:
+                features[layers[name]] = x
+        return features
+
+    def componet_enc(self, x):
+        # x [16,3,256,256]
+        # factor_img [16,7,256,256]
+        vgg_aux = self.get_features(x, self.vgg)  # x是3通道灰度图
+        #x = torch.cat([x, factor_img], dim=1)  # [16,3+7,256,256]
+        x = self.conv1(x) # 64 256 256
+        x = torch.cat([x, vgg_aux['conv1_1']], dim=1)  # 128 256 256
+        x = self.conv2(x)  #  128 128 128
+        x = torch.cat([x, vgg_aux['conv2_1']], dim=1)  # 256 128 128
+        x = self.conv3(x)  # 256 64 64
+        x = torch.cat([x, vgg_aux['conv3_1']], dim=1)  # 512 64 64
+        x = self.conv4(x)  # 512 32 32
+        x = torch.cat([x, vgg_aux['conv4_1']], dim=1)  # 1024 32 32
+        x = self.model(x)
+        return x
+
+    def forward(self, x):
+        code = self.componet_enc(x)
+        return code
+
+class RGBEncoder(nn.Module):
+    def __init__(self, input_dim, dim, n_downsample, n_res, norm, activ, pad_type):
+        super(RGBEncoder, self).__init__()
+        self.model = []
+        self.model += [ConvBlock(input_dim, dim, 7, 1, 3, norm=norm, activation=activ, pad_type=pad_type)]
+        # downsampling blocks
+        for i in range(n_downsample):
+            self.model += [ConvBlock(dim, 2 * dim, 4, 2, 1, norm=norm, activation=activ, pad_type=pad_type)]
+            dim *= 2
+        # residual blocks
+        self.model += [ResBlocks(n_res, dim, norm=norm, activation=activ, pad_type=pad_type)]
+        self.model = nn.Sequential(*self.model)
+        self.output_dim = dim
+
+    def forward(self, x):
+        return self.model(x)
+
+
+class RGBDecoder(nn.Module):
+    def __init__(self, dim, output_dim, n_upsample, n_res, res_norm, activ='relu', pad_type='zero'):
+        super(RGBDecoder, self).__init__()
+        # self.model = []
+        # # AdaIN residual blocks
+        # self.model += [ResBlocks(n_res, dim, res_norm, activ, pad_type=pad_type)]
+        # # upsampling blocks
+        # for i in range(n_upsample):
+        #     self.model += [nn.Upsample(scale_factor=2, mode='nearest'),
+        #                    ConvBlock(dim, dim // 2, 5, 1, 2, norm='ln', activation=activ, pad_type=pad_type)]
+        #     dim //= 2
+        # # use reflection padding in the last conv layer
+        # self.model += [ConvBlock(dim, output_dim, 7, 1, 3, norm='none', activation='tanh', pad_type=pad_type)]
+        # self.model = nn.Sequential(*self.model)
+        #self.Res_Blocks = ModulationResBlocks(n_res, dim, res_norm, activ, pad_type=pad_type)
+        self.mod_conv_1 = ModulationConvBlock(256,256,3)
+        self.mod_conv_2 = ModulationConvBlock(256,256,3)
+        self.mod_conv_3 = ModulationConvBlock(256,256,3)
+        self.mod_conv_4 = ModulationConvBlock(256,256,3)
+        self.mod_conv_5 = ModulationConvBlock(256,256,3)
+        self.mod_conv_6 = ModulationConvBlock(256,256,3)
+        self.mod_conv_7 = ModulationConvBlock(256,256,3)
+        self.mod_conv_8 = ModulationConvBlock(256,256,3)
+        self.upsample_block1 = nn.Upsample(scale_factor=2, mode='nearest')
+        self.conv_1 = ConvBlock(dim, dim // 2, 5, 1, 2, norm='ln', activation=activ, pad_type=pad_type)
+        dim //= 2
+        self.upsample_block2 = nn.Upsample(scale_factor=2, mode='nearest')
+        self.conv_2 = ConvBlock(dim, dim // 2, 5, 1, 2, norm='ln', activation=activ, pad_type=pad_type)
+        dim //= 2
+        self.conv_3 = ConvBlock(dim, output_dim, 7, 1, 3, norm='none', activation='tanh', pad_type=pad_type)
+
+    # def forward(self, x):
+    #     residual = x
+    #     out = self.model(x)
+    #     out += residual
+    #     return out
+    def forward(self, x, code):
+        residual = x
+        x = self.mod_conv_1(x, code[:, :256])
+        x = self.mod_conv_2(x, code[:, 256*1:256*2])
+        x += residual
+        residual = x
+        x = self.mod_conv_2(x, code[:, 256*2:256 * 3])
+        x = self.mod_conv_2(x, code[:, 256*3:256 * 4])
+        x += residual
+        residual =x
+        x = self.mod_conv_2(x, code[:, 256*4:256 * 5])
+        x = self.mod_conv_2(x, code[:, 256*5:256 * 6])
+        x += residual
+        residual = x
+        x = self.mod_conv_2(x, code[:, 256*6:256 * 7])
+        x = self.mod_conv_2(x, code[:, 256*7:256 * 8])
+        x += residual
+        # print(x.shape)
+        x = self.upsample_block1(x)
+        # print(x.shape)
+        x = self.conv_1(x)
+        # print(x_small.shape)
+        x = self.upsample_block2(x)
+        # print(x.shape)
+        x = self.conv_2(x)
+        # print(x_middle.shape)
+        x = self.conv_3(x)
+        # print(x_big.shape)
+        return x
+

models/networks.py

@@ -0,0 +1,244 @@
+import torch
+import torch.nn as nn
+from torch.nn import init
+import functools
+from torch.optim import lr_scheduler
+from .c2pGen import *
+from .p2cGen import *
+from .c2pDis import *
+
+class Identity(nn.Module):
+    def forward(self, x):
+        return x
+
+def get_norm_layer(norm_type='instance'):
+    """Return a normalization layer
+
+    Parameters:
+        norm_type (str) -- the name of the normalization layer: batch | instance | none
+
+    For BatchNorm, we use learnable affine parameters and track running statistics (mean/stddev).
+    For InstanceNorm, we do not use learnable affine parameters. We do not track running statistics.
+    """
+    if norm_type == 'batch':
+        norm_layer = functools.partial(nn.BatchNorm2d, affine=True, track_running_stats=True)
+    elif norm_type == 'instance':
+        norm_layer = functools.partial(nn.InstanceNorm2d, affine=False, track_running_stats=False)
+    elif norm_type == 'none':
+        def norm_layer(x): return Identity()
+    else:
+        raise NotImplementedError('normalization layer [%s] is not found' % norm_type)
+    return norm_layer
+
+
+def get_scheduler(optimizer, opt):
+    """Return a learning rate scheduler
+
+    Parameters:
+        optimizer          -- the optimizer of the network
+        opt (option class) -- stores all the experiment flags; needs to be a subclass of BaseOptions．　
+                              opt.lr_policy is the name of learning rate policy: linear | step | plateau | cosine
+
+    For 'linear', we keep the same learning rate for the first <opt.n_epochs> epochs
+    and linearly decay the rate to zero over the next <opt.n_epochs_decay> epochs.
+    For other schedulers (step, plateau, and cosine), we use the default PyTorch schedulers.
+    See https://pytorch.org/docs/stable/optim.html for more details.
+    """
+    if opt.lr_policy == 'linear':
+        def lambda_rule(epoch):
+            lr_l = 1.0 - max(0, epoch + opt.epoch_count - opt.n_epochs) / float(opt.n_epochs_decay + 1)
+            return lr_l
+        scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda_rule)
+    elif opt.lr_policy == 'step':
+        scheduler = lr_scheduler.StepLR(optimizer, step_size=opt.lr_decay_iters, gamma=0.1)
+    elif opt.lr_policy == 'plateau':
+        scheduler = lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.2, threshold=0.01, patience=5)
+    elif opt.lr_policy == 'cosine':
+        scheduler = lr_scheduler.CosineAnnealingLR(optimizer, T_max=opt.n_epochs, eta_min=0)
+    else:
+        return NotImplementedError('learning rate policy [%s] is not implemented', opt.lr_policy)
+    return scheduler
+
+
+def init_weights(net, init_type='normal', init_gain=0.02):
+    """Initialize network weights.
+
+    Parameters:
+        net (network)   -- network to be initialized
+        init_type (str) -- the name of an initialization method: normal | xavier | kaiming | orthogonal
+        init_gain (float)    -- scaling factor for normal, xavier and orthogonal.
+
+    """
+    def init_func(m):  # define the initialization function
+        classname = m.__class__.__name__
+        if hasattr(m, 'weight') and (classname.find('Conv') != -1 or classname.find('Linear') != -1):
+            if init_type == 'normal':
+                init.normal_(m.weight.data, 0.0, init_gain)
+            elif init_type == 'xavier':
+                init.xavier_normal_(m.weight.data, gain=init_gain)
+            elif init_type == 'kaiming':
+                init.kaiming_normal_(m.weight.data, a=0, mode='fan_in')
+            elif init_type == 'orthogonal':
+                init.orthogonal_(m.weight.data, gain=init_gain)
+            else:
+                raise NotImplementedError('initialization method [%s] is not implemented' % init_type)
+            if hasattr(m, 'bias') and m.bias is not None:
+                init.constant_(m.bias.data, 0.0)
+        elif classname.find('BatchNorm2d') != -1:  # BatchNorm Layer's weight is not a matrix; only normal distribution applies.
+            init.normal_(m.weight.data, 1.0, init_gain)
+            init.constant_(m.bias.data, 0.0)
+
+    #print('initialize network with %s' % init_type)
+    net.apply(init_func)  # apply the initialization function <init_func>
+
+
+def init_net(net, init_type='normal', init_gain=0.02, gpu_ids=[]):
+    """Initialize a network: 1. register CPU/GPU device (with multi-GPU support); 2. initialize the network weights
+    Parameters:
+        net (network)      -- the network to be initialized
+        init_type (str)    -- the name of an initialization method: normal | xavier | kaiming | orthogonal
+        gain (float)       -- scaling factor for normal, xavier and orthogonal.
+        gpu_ids (int list) -- which GPUs the network runs on: e.g., 0,1,2
+
+    Return an initialized network.
+    """
+    gpu_ids = [0]
+    if len(gpu_ids) > 0:
+        # assert(torch.cuda.is_available()) #uncomment this for using gpu
+        net.to(torch.device("cpu")) #change this for using gpu to gpu_ids[0]
+        net = torch.nn.DataParallel(net, gpu_ids)  # multi-GPUs
+    init_weights(net, init_type, init_gain=init_gain)
+    return net
+
+
+def define_G(input_nc, output_nc, ngf, netG, norm='batch', use_dropout=False, init_type='normal', init_gain=0.02, gpu_ids=[]):
+    """Create a generator
+
+    Parameters:
+        input_nc (int) -- the number of channels in input images
+        output_nc (int) -- the number of channels in output images
+        ngf (int) -- the number of filters in the last conv layer
+        netG (str) -- the architecture's name: resnet_9blocks | resnet_6blocks | unet_256 | unet_128
+        norm (str) -- the name of normalization layers used in the network: batch | instance | none
+        use_dropout (bool) -- if use dropout layers.
+        init_type (str)    -- the name of our initialization method.
+        init_gain (float)  -- scaling factor for normal, xavier and orthogonal.
+        gpu_ids (int list) -- which GPUs the network runs on: e.g., 0,1,2
+
+    Returns a generator
+    """
+    net = None
+    norm_layer = get_norm_layer(norm_type=norm)
+
+    if netG == 'c2pGen':  #                     style_dim  mlp_dim
+        net = C2PGen(input_nc, output_nc, ngf, 2, 4, 256, 256, activ='relu', pad_type='reflect')
+        #print('c2pgen resblock is 8')
+    elif netG == 'p2cGen':
+        net = P2CGen(input_nc, output_nc, ngf, 2, 3, activ='relu', pad_type='reflect')
+    elif netG == 'antialias':
+        net = AliasNet(input_nc, output_nc, ngf, 2, 3, activ='relu', pad_type='reflect')
+    else:
+        raise NotImplementedError('Generator model name [%s] is not recognized' % netG)
+    return init_net(net, init_type, init_gain, gpu_ids)
+
+
+
+def define_D(input_nc, ndf, netD, n_layers_D=3, norm='batch', init_type='normal', init_gain=0.02, gpu_ids=[]):
+    """Create a discriminator
+
+    Parameters:
+        input_nc (int)     -- the number of channels in input images
+        ndf (int)          -- the number of filters in the first conv layer
+        netD (str)         -- the architecture's name: basic | n_layers | pixel
+        n_layers_D (int)   -- the number of conv layers in the discriminator; effective when netD=='n_layers'
+        norm (str)         -- the type of normalization layers used in the network.
+        init_type (str)    -- the name of the initialization method.
+        init_gain (float)  -- scaling factor for normal, xavier and orthogonal.
+        gpu_ids (int list) -- which GPUs the network runs on: e.g., 0,1,2
+
+    Returns a discriminator
+    """
+    net = None
+    norm_layer = get_norm_layer(norm_type=norm)
+
+
+    if netD == 'CPDis':
+        net = CPDis(image_size=256, conv_dim=64, repeat_num=3, norm='SN')
+    elif netD == 'CPDis_cls':
+        net = CPDis_cls(image_size=256, conv_dim=64, repeat_num=3, norm='SN')
+    else:
+        raise NotImplementedError('Discriminator model name [%s] is not recognized' % netD)
+    return init_net(net, init_type, init_gain, gpu_ids)
+
+
+class GANLoss(nn.Module):
+    """Define different GAN objectives.
+
+    The GANLoss class abstracts away the need to create the target label tensor
+    that has the same size as the input.
+    """
+
+    def __init__(self, gan_mode, target_real_label=1.0, target_fake_label=0.0):
+        """ Initialize the GANLoss class.
+
+        Parameters:
+            gan_mode (str) - - the type of GAN objective. It currently supports vanilla, lsgan, and wgangp.
+            target_real_label (bool) - - label for a real image
+            target_fake_label (bool) - - label of a fake image
+
+        Note: Do not use sigmoid as the last layer of Discriminator.
+        LSGAN needs no sigmoid. vanilla GANs will handle it with BCEWithLogitsLoss.
+        """
+        super(GANLoss, self).__init__()
+        self.register_buffer('real_label', torch.tensor(target_real_label))
+        self.register_buffer('fake_label', torch.tensor(target_fake_label))
+        self.gan_mode = gan_mode
+        if gan_mode == 'lsgan':
+            self.loss = nn.MSELoss()
+        elif gan_mode == 'vanilla':
+            self.loss = nn.BCEWithLogitsLoss()
+        elif gan_mode in ['wgangp']:
+            self.loss = None
+        else:
+            raise NotImplementedError('gan mode %s not implemented' % gan_mode)
+
+    def get_target_tensor(self, prediction, target_is_real):
+        """Create label tensors with the same size as the input.
+
+        Parameters:
+            prediction (tensor) - - tpyically the prediction from a discriminator
+            target_is_real (bool) - - if the ground truth label is for real images or fake images
+
+        Returns:
+            A label tensor filled with ground truth label, and with the size of the input
+        """
+
+        if target_is_real:
+            target_tensor = self.real_label
+        else:
+            target_tensor = self.fake_label
+        return target_tensor.expand_as(prediction)
+
+    def __call__(self, prediction, target_is_real):
+        """Calculate loss given Discriminator's output and grount truth labels.
+
+        Parameters:
+            prediction (tensor) - - tpyically the prediction output from a discriminator
+            target_is_real (bool) - - if the ground truth label is for real images or fake images
+
+        Returns:
+            the calculated loss.
+        """
+        if self.gan_mode in ['lsgan', 'vanilla']:
+            target_tensor = self.get_target_tensor(prediction, target_is_real)
+            loss = self.loss(prediction, target_tensor)
+        elif self.gan_mode == 'wgangp':
+            if target_is_real:
+                loss = -prediction.mean()
+            else:
+                loss = prediction.mean()
+        return loss
+
+
+
+

models/p2cGen.py

@@ -0,0 +1,76 @@
+from .basic_layer import *
+
+
+class P2CGen(nn.Module):
+    def __init__(self, input_dim, output_dim, dim, n_downsample, n_res, activ='relu', pad_type='reflect'):
+        super(P2CGen, self).__init__()
+        self.RGBEnc = RGBEncoder(input_dim, dim, n_downsample, n_res, "in", activ, pad_type=pad_type)
+        self.RGBDec = RGBDecoder(self.RGBEnc.output_dim, output_dim, n_downsample, n_res, res_norm='in',
+                                      activ=activ, pad_type=pad_type)
+
+    def forward(self, x):
+        x = self.RGBEnc(x)
+        # print("encoder->>", x.shape)
+        x = self.RGBDec(x)
+        # print(x_small.shape)
+        # print(x_middle.shape)
+        # print(x_big.shape)
+        #return y_small, y_middle, y_big
+        return x
+
+
+class RGBEncoder(nn.Module):
+    def __init__(self, input_dim, dim, n_downsample, n_res, norm, activ, pad_type):
+        super(RGBEncoder, self).__init__()
+        self.model = []
+        self.model += [ConvBlock(input_dim, dim, 7, 1, 3, norm=norm, activation=activ, pad_type=pad_type)]
+        # downsampling blocks
+        for i in range(n_downsample):
+            self.model += [ConvBlock(dim, 2 * dim, 4, 2, 1, norm=norm, activation=activ, pad_type=pad_type)]
+            dim *= 2
+        # residual blocks
+        self.model += [ResBlocks(n_res, dim, norm=norm, activation=activ, pad_type=pad_type)]
+        self.model = nn.Sequential(*self.model)
+        self.output_dim = dim
+
+    def forward(self, x):
+        return self.model(x)
+
+
+class RGBDecoder(nn.Module):
+    def __init__(self, dim, output_dim, n_upsample, n_res, res_norm, activ='relu', pad_type='zero'):
+        super(RGBDecoder, self).__init__()
+        # self.model = []
+        # # AdaIN residual blocks
+        # self.model += [ResBlocks(n_res, dim, res_norm, activ, pad_type=pad_type)]
+        # # upsampling blocks
+        # for i in range(n_upsample):
+        #     self.model += [nn.Upsample(scale_factor=2, mode='nearest'),
+        #                    ConvBlock(dim, dim // 2, 5, 1, 2, norm='ln', activation=activ, pad_type=pad_type)]
+        #     dim //= 2
+        # # use reflection padding in the last conv layer
+        # self.model += [ConvBlock(dim, output_dim, 7, 1, 3, norm='none', activation='tanh', pad_type=pad_type)]
+        # self.model = nn.Sequential(*self.model)
+        self.Res_Blocks = ResBlocks(n_res, dim, res_norm, activ, pad_type=pad_type)
+        self.upsample_block1 = nn.Upsample(scale_factor=2, mode='nearest')
+        self.conv_1 = ConvBlock(dim, dim // 2, 5, 1, 2, norm='ln', activation=activ, pad_type=pad_type)
+        dim //= 2
+        self.upsample_block2 = nn.Upsample(scale_factor=2, mode='nearest')
+        self.conv_2 = ConvBlock(dim, dim // 2, 5, 1, 2, norm='ln', activation=activ, pad_type=pad_type)
+        dim //= 2
+        self.conv_3 = ConvBlock(dim, output_dim, 7, 1, 3, norm='none', activation='tanh', pad_type=pad_type)
+
+    def forward(self, x):
+        x = self.Res_Blocks(x)
+        # print(x.shape)
+        x = self.upsample_block1(x)
+        # print(x.shape)
+        x = self.conv_1(x)
+        # print(x_small.shape)
+        x = self.upsample_block2(x)
+        # print(x.shape)
+        x = self.conv_2(x)
+        # print(x_middle.shape)
+        x = self.conv_3(x)
+        # print(x_big.shape)
+        return x

pixelization.py

@@ -0,0 +1,151 @@
+import os
+import torch
+import torchvision.transforms as transforms
+from PIL import Image
+import numpy as np
+from models.networks import define_G
+import glob
+
+
+class Model():
+    def __init__(self, device="cpu"):
+        self.device = torch.device(device)
+        self.G_A_net = None
+        self.alias_net = None
+        self.ref_t = None
+
+    def load(self):
+        with torch.no_grad():
+            self.G_A_net = define_G(3, 3, 64, "c2pGen", "instance", False, "normal", 0.02, [0])
+            self.alias_net = define_G(3, 3, 64, "antialias", "instance", False, "normal", 0.02, [0])
+
+            G_A_state = torch.load("160_net_G_A.pth" if not os.environ['NET_MODEL'] else os.environ['NET_MODEL'], map_location=str(self.device))
+            for p in list(G_A_state.keys()):
+                G_A_state["module."+str(p)] = G_A_state.pop(p)
+            self.G_A_net.load_state_dict(G_A_state)
+
+            alias_state = torch.load("alias_net.pth" if not os.environ['ALIAS_MODEL'] else os.environ['ALIAS_MODEL'], map_location=str(self.device))
+            for p in list(alias_state.keys()):
+                alias_state["module."+str(p)] = alias_state.pop(p)
+            self.alias_net.load_state_dict(alias_state)
+
+            ref_img = Image.open("reference.png").convert('L')
+            self.ref_t = process(greyscale(ref_img)).to(self.device)
+
+    def pixelize(self, in_img, out_img):
+        with torch.no_grad():
+            in_img = Image.open(in_img).convert('RGB')
+            in_t = process(in_img).to(self.device)
+        
+            out_t = self.alias_net(self.G_A_net(in_t, self.ref_t))
+
+            save(out_t, out_img)
+
+    def pixelize_modified(self, in_img, pixel_size, upscale_after) -> Image.Image:
+        with torch.no_grad():
+            in_img = in_img.convert('RGB')
+
+            # limit in_img size to 1024x1024 so it didn't destroyed by large image
+            if in_img.size[0] > 1024 or in_img.size[1] > 1024:
+                in_img.thumbnail((1024, 1024), Image.NEAREST)
+
+            in_img.resize((in_img.size[0] * 4 // pixel_size, in_img.size[1] * 4 // pixel_size))
+
+            in_t = process(in_img).to(self.device)
+
+            out_t = self.alias_net(self.G_A_net(in_t, self.ref_t))
+            img = to_image(out_t, pixel_size, upscale_after)
+            return img
+
+def to_image(tensor, pixel_size, upscale_after):
+    img = tensor.data[0].cpu().float().numpy()
+    img = (np.transpose(img, (1, 2, 0)) + 1) / 2.0 * 255.0
+    img = img.astype(np.uint8)
+    img = Image.fromarray(img)
+    img = img.resize((img.size[0]//4, img.size[1]//4), resample=Image.Resampling.NEAREST)
+    if upscale_after:
+        img = img.resize((img.size[0]*pixel_size, img.size[1]*pixel_size), resample=Image.Resampling.NEAREST)
+
+    return img
+
+
+def greyscale(img):
+    gray = np.array(img.convert('L'))
+    tmp = np.expand_dims(gray, axis=2)
+    tmp = np.concatenate((tmp, tmp, tmp), axis=-1)
+    return Image.fromarray(tmp)
+
+def process(img):
+    ow,oh = img.size
+
+    nw = int(round(ow / 4) * 4)
+    nh = int(round(oh / 4) * 4)
+
+    left = (ow - nw)//2
+    top = (oh - nh)//2
+    right = left + nw
+    bottom = top + nh
+
+    img = img.crop((left, top, right, bottom))
+
+    trans = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
+
+    return trans(img)[None, :, :, :]
+
+def save(tensor, file):
+    img = tensor.data[0].cpu().float().numpy()
+    img = (np.transpose(img, (1, 2, 0)) + 1) / 2.0 * 255.0
+    img = img.astype(np.uint8)
+    img = Image.fromarray(img)
+    img = img.resize((img.size[0]//4, img.size[1]//4), resample=Image.Resampling.NEAREST)
+    img = img.resize((img.size[0]*4, img.size[1]*4), resample=Image.Resampling.NEAREST)
+    img.save(file)
+
+def pixelize_cli():
+    import argparse
+    import os
+    parser = argparse.ArgumentParser(description='Pixelization')
+    parser.add_argument('--input', type=str, default=None, required=True, help='path to image or directory')
+    parser.add_argument('--output', type=str, default=None, required=False, help='path to save image/images')
+    parser.add_argument('--cpu', action='store_true', help='use CPU instead of GPU')
+
+    args = parser.parse_args()
+    in_path = args.input
+    out_path = args.output
+    use_cpu = args.cpu
+
+    if not os.path.exists("alias_net.pth" if not os.environ['ALIAS_MODEL'] else os.environ['ALIAS_MODEL']):
+        print("missing models")
+
+    pairs = []
+
+    if os.path.isdir(in_path):
+        in_images = glob.glob(in_path + "/*.png") + glob.glob(in_path + "/*.jpg")
+        if not out_path:
+            out_path = os.path.join(in_path, "outputs")
+        if not os.path.exists(out_path):
+            os.makedirs(out_path)
+        elif os.path.isfile(out_path):
+            print("output cant be a file if input is a directory")
+            return
+        for i in in_images:
+            pairs += [(i, i.replace(in_path, out_path))]
+    elif os.path.isfile(in_path):
+        if not out_path:
+            base, ext = os.path.splitext(in_path)
+            out_path = base+"_pixelized"+ext
+        else:
+            if os.path.isdir(out_path):
+                _, file = os.path.split(in_path)
+                out_path = os.path.join(out_path, file)
+        pairs = [(in_path, out_path)]
+    
+    m = Model(device = "cpu" if use_cpu else "cuda")
+    m.load()
+    
+    for in_file, out_file in pairs:
+        print("PROCESSING", in_file, "TO", out_file)
+        m.pixelize(in_file, out_file)
+
+if __name__ == "__main__":
+    pixelize_cli()

requirements.txt

@@ -0,0 +1,5 @@
+torch
+torchvision
+transforms
+numpy==1.24.1
+pillow