update

Author: SerialLain3170

GauGAN/dataset.py

@@ -4,33 +4,49 @@
 import copy
 import chainer
 
-from xdog import line_process
+from xdog import xdog_process
 from chainer import cuda
 
 xp = cuda.cupy
 cuda.get_device(0).use()
 
 
 class DataLoader:
-    def __init__(self, path):
+    def __init__(self,
+                 path,
+                 extension='.jpg',
+                 img_size=224,
+                 latent_dim=256):
+
         self.path = path
-        self.pathlist = list(self.path.glob('**/*.jpg'))
+        self.pathlist = list(self.path.glob(f"**/*{extension}"))
         self.train, self.valid = self._split(self.pathlist)
         self.train_len = len(self.train)
         self.valid_len = len(self.valid)
 
+        self.size = img_size
+        self.latent_dim = latent_dim
+
+        self.interpolations = (
+            cv.INTER_LINEAR,
+            cv.INTER_AREA,
+            cv.INTER_NEAREST,
+            cv.INTER_CUBIC,
+            cv.INTER_LANCZOS4
+        )
+
     def __str__(self):
         return f"dataset path: {self.path} train data: {self.train_len}"
 
     def _split(self, pathlist: list):
-        split_point = int(len(self.pathlist) * 0.9)
+        split_point = int(len(self.pathlist) * 0.95)
         x_train = self.pathlist[:split_point]
         x_test = self.pathlist[split_point:]
 
         return x_train, x_test
 
     @staticmethod
-    def _random_crop(line, color, size=224):
+    def _random_crop(line, color, size):
         height, width = line.shape[0], line.shape[1]
         rnd0 = np.random.randint(height - size - 1)
         rnd1 = np.random.randint(width - size - 1)
@@ -52,76 +68,27 @@ def _coordinate(image):
     def _variable(image_list):
         return chainer.as_variable(xp.array(image_list).astype(xp.float32))
 
-    @staticmethod
-    def noise_generator(batchsize):
-        noise = xp.random.normal(size=(batchsize, 256)).astype(xp.float32)
+    def noise_generator(self, batchsize):
+        noise = xp.random.normal(size=(batchsize, self.latent_dim)).astype(xp.float32)
 
         return chainer.as_variable(noise)
 
-    @staticmethod
-    def _making_mask(mask, color, size=224):
-        choice = np.random.choice(['width', 'height', 'diag'])
-
-        if choice == 'width':
-            rnd_height = np.random.randint(4, 8)
-            rnd_width = np.random.randint(4, 64)
-
-            rnd1 = np.random.randint(size - rnd_height)
-            rnd2 = np.random.randint(size - rnd_width)
-            mask[rnd1:rnd1+rnd_height, rnd2:rnd2+rnd_width] = color[rnd1:rnd1+rnd_height, rnd2:rnd2+rnd_width]
-
-        elif choice == 'height':
-            rnd_height = np.random.randint(4, 64)
-            rnd_width = np.random.randint(4, 8)
-
-            rnd1 = np.random.randint(size - rnd_height)
-            rnd2 = np.random.randint(size - rnd_width)
-            mask[rnd1:rnd1+rnd_height, rnd2:rnd2+rnd_width] = color[rnd1:rnd1+rnd_height, rnd2:rnd2+rnd_width]
-
-        elif choice == 'diag':
-            rnd_height = np.random.randint(4, 8)
-            rnd_width = np.random.randint(4, 64)
-
-            rnd1 = np.random.randint(size - rnd_height - rnd_width - 1)
-            rnd2 = np.random.randint(size - rnd_width)
-
-            for index in range(rnd_width):
-                mask[rnd1 + index : rnd1 + rnd_height + index, rnd2 + index] = color[rnd1 + index: rnd1 + rnd_height + index, rnd2 + index]
-
-        return mask
-
-    def _prepare_pair(self, image_path, size=224, repeat=16):
-        interpolations = (
-            cv.INTER_LINEAR,
-            cv.INTER_AREA,
-            cv.INTER_NEAREST,
-            cv.INTER_CUBIC,
-            cv.INTER_LANCZOS4
-        )
-        interpolation = random.choice(interpolations)
+    def _prepare_pair(self, image_path, size, repeat=16):
+        interpolation = random.choice(self.interpolations)
 
         color = cv.imread(str(image_path))
-        line = line_process(str(image_path))
+        line = xdog_process(str(image_path))
 
         line, color = self._random_crop(line, color, size=size)
-        mask = copy.copy(line)
-
-        for _ in range(repeat):
-            mask = self._making_mask(mask, color, size=size)
-        mask_ds = cv.resize(mask, (int(size/2), int(size/2)), interpolation=interpolation)
 
         color = self._coordinate(color)
         line = self._coordinate(line)
-        mask = self._coordinate(mask)
-        mask_ds = self._coordinate(mask_ds)
 
-        return (color, line, mask, mask_ds)
+        return (color, line)
 
-    def __call__(self, batchsize, mode='train', size=224):
+    def __call__(self, batchsize, mode='train'):
         color_box = []
         line_box = []
-        mask_box = []
-        mask_ds_box = []
 
         for _ in range(batchsize):
             if mode == 'train':
@@ -133,16 +100,12 @@ def __call__(self, batchsize, mode='train', size=224):
             else:
                 raise AttributeError
 
-            color, line, mask, mask_ds = self._prepare_pair(image_path, size=size)
+            color, line = self._prepare_pair(image_path, size=self.size)
 
             color_box.append(color)
             line_box.append(line)
-            mask_box.append(mask)
-            mask_ds_box.append(mask_ds)
 
         color = self._variable(color_box)
         line = self._variable(line_box)
-        mask = self._variable(mask_box)
-        mask_ds = self._variable(mask_ds_box)
 
-        return (color, line, mask, mask_ds)
+        return (color, line)

GauGAN/evaluation.py

@@ -18,19 +18,21 @@ def __call__(self, y, t, x, outdir, epoch, validsize=3):
         pylab.rcParams['figure.figsize'] = (16.0, 16.0)
         pylab.clf()
 
+        wid = int(validsize/2)
+
         for index in range(validsize):
             tmp = self._coordinate(x[index])
-            pylab.subplot(validsize, validsize, validsize * index + 1)
+            pylab.subplot(wid, wid, 3 * index + 1)
             pylab.imshow(tmp)
             pylab.axis('off')
             pylab.savefig(f"{outdir}/visualize_{epoch}.png")
             tmp = self._coordinate(t[index])
-            pylab.subplot(validsize, validsize, validsize * index + 2)
+            pylab.subplot(wid, wid, 3 * index + 2)
             pylab.imshow(tmp)
             pylab.axis('off')
             pylab.savefig(f"{outdir}/visualize_{epoch}.png")
             tmp = self._coordinate(y[index])
-            pylab.subplot(validsize, validsize, validsize * index + 3)
+            pylab.subplot(wid, wid, 3 * index + 3)
             pylab.imshow(tmp)
             pylab.axis('off')
             pylab.savefig(f"{outdir}/visualize_{epoch}.png")

GauGAN/model.py

@@ -121,10 +121,10 @@ def __call__(self, x, c):
         return F.tanh(h)
 
 
-class Discriminator(Chain):
+class DiscriminatorBlock(Chain):
     def __init__(self, base=64):
         w = initializers.GlorotUniform()
-        super(Discriminator, self).__init__()
+        super(DiscriminatorBlock, self).__init__()
         with self.init_scope():
             self.c0 = SNConvolution2D(6, base, 4, 2, 1, initialW=w)
             self.c1 = SNConvolution2D(base, base*2, 4, 2, 1, initialW=w)
@@ -148,6 +148,30 @@ def __call__(self, x):
         return h, [h1, h2, h3, h4]
 
 
+class Discriminator(Chain):
+    def __init__(self, base=64):
+        super(Discriminator, self).__init__()
+        discriminators = chainer.ChainList()
+        for _ in range(3):
+            discriminators.add_link(DiscriminatorBlock())
+        with self.init_scope():
+            self.dis = discriminators
+
+    def __call__(self, x):
+        adv_list = []
+        feat_list = []
+
+        for index in range(3):
+            h, h_list = self.dis[index](x)
+
+            adv_list.append(h)
+            feat_list.append(h_list)
+
+            x = F.average_pooling_2d(x, 3, 2, 1)
+
+        return adv_list, feat_list
+
+
 class Prior(chainer.Link):
 
     def __init__(self):
@@ -157,4 +181,4 @@ def __init__(self):
         self.scale = xp.ones(256, xp.float32)
 
     def __call__(self):
-        return D.Normal(self.loc, scale=self.scale)
+        return D.Normal(self.loc, scale=self.scale)

GauGAN/sn.py

@@ -8,6 +8,7 @@
 from chainer.links.connection.linear import Linear
 import chainer.functions as F
 
+
 def _l2normalize(v, eps=1e-12):
     norm = cuda.reduce('T x', 'T out',
                        'x * x', 'a + b', 'out = sqrt(a)', 0,
@@ -19,6 +20,7 @@ def _l2normalize(v, eps=1e-12):
 
     return div(v, norm(v), eps)
 
+
 def max_singular_value(W, u=None, Ip=1):
     """
     Apply power iteration for the weight parameter
@@ -37,6 +39,7 @@ def max_singular_value(W, u=None, Ip=1):
 
     return sigma, _u, _v
 
+
 class SNConvolution2D(Convolution2D):
     """Two-dimensional convolutional layer with spectral normalization.
     This link wraps the :func:`~chainer.functions.convolution_2d` function and
@@ -129,6 +132,7 @@ def __call__(self, x):
         return convolution_2d.convolution_2d(
             x, self.W_bar, self.b, self.stride, self.pad)
 
+
 class SNLinear(Linear):
     """Linear layer with Spectral Normalization.
     Args:
@@ -204,4 +208,4 @@ def __call__(self, x):
         """
         if self.W.data is None:
             self._initialize_params(x.size // x.shape[0])
-        return linear.linear(x, self.W_bar, self.b)
+        return linear.linear(x, self.W_bar, self.b)

GauGAN/train.py

@@ -14,47 +14,67 @@
 cuda.get_device(0).use()
 
 
-def downsampling(array):
-    d2 = F.average_pooling_2d(array, 3, 2, 1)
-    d4 = F.average_pooling_2d(d2, 3, 2, 1)
-
-    return d2, d4
-
-
 class GauGANLossFunction:
     def __init__(self):
         pass
 
     @staticmethod
     def content_loss(y, t):
-        return F.mean_absolute_error(y, t)
+        return 10.0 * F.mean_absolute_error(y, t)
 
     @staticmethod
-    def dis_hinge_loss(discriminator, y, t):
-        y_dis, _ = discriminator(y)
-        t_dis, _ = discriminator(t)
+    def dis_loss(discriminator, y, t):
+        y_adv_list, _ = discriminator(y)
+        t_adv_list, _ = discriminator(t)
+
+        sum_loss = 0
 
-        return F.mean(F.relu(1. - t_dis)) + F.mean(F.relu(1. + y_dis))
+        for y_adv, t_adv in zip(y_adv_list, t_adv_list):
+            loss = F.mean(F.relu(1. - t_adv)) + F.mean(F.relu(1. + y_adv))
+            sum_loss += loss
+
+        return sum_loss
 
     @staticmethod
-    def gen_hinge_loss(discriminator, y, t):
-        y_dis, y_feats = discriminator(y)
+    def gen_loss(discriminator, y, t):
+        y_dis_list, y_feats = discriminator(y)
         _, t_feats = discriminator(t)
 
         sum_loss = 0
-        for yf, tf in zip(y_feats, t_feats):
-            _, ch, height, width = yf.shape
-            sum_loss += 10.0 * F.mean_absolute_error(yf, tf) / (ch * height * width)
-
-        return -F.mean(y_dis) + sum_loss
 
+        # adversarial loss
+        for y_dis in y_dis_list:
+            loss = -F.mean(y_dis)
+            sum_loss += loss
+
+        # feature matching loss
+        for yf_list, tf_list in zip(y_feats, t_feats):
+            for yf, tf in zip(yf_list, tf_list):
+                _, ch, height, width = yf.shape
+                sum_loss += 10.0 * F.mean_absolute_error(yf, tf) / (ch * height * width)
+
+        return sum_loss
+
+
+def train(epochs,
+          iterations,
+          batchsize,
+          validsize,
+          outdir,
+          modeldir,
+          data_path,
+          extension,
+          img_size,
+          latent_dim,
+          learning_rate,
+          beta1,
+          beta2,
+          enable):
 
-def train(epochs, iterations, batchsize, validsize, path, outdir,
-          con_weight, kl_weight, enable):
     # Dataset Definition
-    dataloader = DataLoader(path)
+    dataloader = DataLoader(data_path, extension, img_size, latent_dim)
     print(dataloader)
-    color_valid, line_valid, _, _ = dataloader(validsize, mode="valid")
+    color_valid, line_valid = dataloader(validsize, mode="valid")
     noise_valid = dataloader.noise_generator(validsize)
 
     # Model Definition
@@ -65,19 +85,11 @@ def train(epochs, iterations, batchsize, validsize, path, outdir,
 
     generator = Generator()
     generator.to_gpu()
-    gen_opt = set_optimizer(generator)
+    gen_opt = set_optimizer(generator, learning_rate, beta1, beta2)
 
     discriminator = Discriminator()
     discriminator.to_gpu()
-    dis_opt = set_optimizer(discriminator)
-
-    discriminator_d2 = Discriminator()
-    discriminator_d2.to_gpu()
-    dis2_opt = set_optimizer(discriminator_d2)
-
-    discriminator_d4 = Discriminator()
-    discriminator_d4.to_gpu()
-    dis4_opt = set_optimizer(discriminator_d4)
+    dis_opt = set_optimizer(discriminator, learning_rate, beta1, beta2)
 
     # Loss Funtion Definition
     lossfunc = GauGANLossFunction()
@@ -86,95 +98,64 @@ def train(epochs, iterations, batchsize, validsize, path, outdir,
     evaluator = Evaluaton()
 
     for epoch in range(epochs):
-        sum_loss = 0
+        sum_dis_loss = 0
+        sum_gen_loss = 0
         for batch in range(0, iterations, batchsize):
-            color, line, _, _ = dataloader(batchsize)
-
-            color_d2, color_d4 = downsampling(color)
-            line_d2, line_d4 = downsampling(line)
+            color, line = dataloader(batchsize)
             z = dataloader.noise_generator(batchsize)
 
+            # Discriminator update
             if enable:
                 mu, sigma = encoder(color)
                 z = F.gaussian(mu, sigma)
             y = generator(z, line)
-            y_d2, y_d4 = downsampling(y)
 
             y.unchain_backward()
-            y_d2.unchain_backward()
-            y_d4.unchain_backward()
 
-            loss = lossfunc.dis_hinge_loss(
+            dis_loss = lossfunc.dis_loss(
                 discriminator,
                 F.concat([y, line]),
                 F.concat([color, line])
             )
-            loss += lossfunc.dis_hinge_loss(
-                discriminator_d2,
-                F.concat([y_d2, line_d2]),
-                F.concat([color_d2, line_d2])
-            )
-            loss += lossfunc.dis_hinge_loss(
-                discriminator_d4,
-                F.concat([y_d4, line_d4]),
-                F.concat([color_d4, line_d4])
-            )
 
             discriminator.cleargrads()
-            discriminator_d2.cleargrads()
-            discriminator_d4.cleargrads()
-            loss.backward()
+            dis_loss.backward()
             dis_opt.update()
-            dis2_opt.update()
-            dis4_opt.update()
-            loss.unchain_backward()
+            dis_loss.unchain_backward()
+
+            sum_dis_loss += dis_loss.data
 
+            # Generator update
             z = dataloader.noise_generator(batchsize)
 
             if enable:
                 mu, sigma = encoder(color)
                 z = F.gaussian(mu, sigma)
             y = generator(z, line)
-            y_d2, y_d4 = downsampling(y)
 
-            loss = lossfunc.gen_hinge_loss(
+            gen_loss = lossfunc.gen_loss(
                 discriminator,
                 F.concat([y, line]),
                 F.concat([color, line])
             )
-            loss += lossfunc.gen_hinge_loss(
-                discriminator_d2,
-                F.concat([y_d2, line_d2]),
-                F.concat([color_d2, line_d2])
-            )
-            loss += lossfunc.gen_hinge_loss(
-                discriminator_d4,
-                F.concat([y_d4, line_d4]),
-                F.concat([color_d4, line_d4])
-            )
-            loss += con_weight * lossfunc.content_loss(y, color)
-            loss += con_weight * lossfunc.content_loss(y_d2, color_d2)
-            loss += con_weight * lossfunc.content_loss(y_d4, color_d4)
+            gen_loss += lossfunc.content_loss(y, color)
 
             if enable:
-                loss += kl_weight * F.gaussian_kl_divergence(mu, sigma) / batchsize
+                gen_loss += 0.05 * F.gaussian_kl_divergence(mu, sigma) / batchsize
 
             generator.cleargrads()
             if enable:
                 encoder.cleargrads()
-            loss.backward()
+            gen_loss.backward()
             gen_opt.update()
             if enable:
                 enc_opt.update()
-            loss.unchain_backward()
+            gen_loss.unchain_backward()
 
-            sum_loss += loss.data
+            sum_gen_loss += gen_loss.data
 
             if batch == 0:
-                serializers.save_npz(f"{outdir}/generator.model", generator)
-                serializers.save_npz(f"{outdir}/discriminator_0.model", discriminator)
-                serializers.save_npz(f"{outdir}/discriminator_2.model", discriminator_d2)
-                serializers.save_npz(f"{outdir}/discriminator_4.model", discriminator_d4)
+                serializers.save_npz(f"{modeldir}/generator_{epoch}.model", generator)
 
                 with chainer.using_config("train", False):
                     y = generator(noise_valid, line_valid)
@@ -183,25 +164,35 @@ def train(epochs, iterations, batchsize, validsize, path, outdir,
                 cr = color_valid.data.get()
 
                 evaluator(y, cr, sr, outdir, epoch, validsize=validsize)
-    
-    print(f"epoch: {epoch}")
-    print(f"loss: {sum_loss / iterations}")
+
+        print(f"epoch: {epoch}")
+        print(f"dis loss: {sum_dis_loss / iterations} gen loss: {sum_gen_loss / iterations}")
 
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(description="GauGAN")
     parser.add_argument('--e', type=int, default=1000, help="the number of epochs")
-    parser.add_argument('--i', type=int, default=10000, help="the number of iterations")
+    parser.add_argument('--i', type=int, default=2000, help="the number of iterations")
     parser.add_argument('--b', type=int, default=16, help="batch size")
-    parser.add_argument('--v', type=int, default=3, help="valid size")
-    parser.add_argument('--w', type=float, default=10.0, help="the weight of content loss")
-    parser.add_argument('--kl', type=float, default=0.05, help="the weight of kl divergence loss")
+    parser.add_argument('--v', type=int, default=12, help="valid size")
+    parser.add_argument('--outdir', type=Path, default='outdir', help="output directory")
+    parser.add_argument('--modeldir', type=Path, default='modeldir', help="model output directory")
+    parser.add_argument('--ext', type=str, default=".jpg", help="extension of training images")
+    parser.add_argument('--size', type=int, default=224, help="the size of training images")
+    parser.add_argument('--dim', type=int, default=256, help="dimensions of latent space")
+    parser.add_argument('--lr', type=float, default=0.0002, help="learning rate of Adam")
+    parser.add_argument('--b1', type=float, default=0.0, help="beta1 of Adam")
+    parser.add_argument('--b2', type=float, default=0.999, help="beta2 of Adam")
+    parser.add_argument('--data_path', type=Path, help="path which contains training data")
     parser.add_argument('--encoder', action="store_true", help="enable image encoder")
 
     args = parser.parse_args()
 
-    dataset_path = Path('./Dataset/danbooru-images/')
-    outdir = Path('./outdir')
+    outdir = args.outdir
     outdir.mkdir(exist_ok=True)
 
-    train(args.e, args.i, args.b, args.v, dataset_path, outdir, args.w, args.kl, args.encoder)
+    modeldir = args.modeldir
+    modeldir.mkdir(exist_ok=True)
+
+    train(args.e, args.i, args.b, args.v, outdir, modeldir, args.data_path,
+          args.ext, args.size, args.dim, args.lr, args.b1, args.b2, args.encoder)

GauGAN/xdog.py

@@ -1,38 +1,43 @@
 import numpy as np
 import cv2 as cv
 
+
 def sharpImage(img, sigma, k_sigma, p):
     sigma_large = sigma * k_sigma
-    G_small = cv.GaussianBlur(img,(0, 0), sigma)
-    G_large = cv.GaussianBlur(img,(0, 0), sigma_large)
+    G_small = cv.GaussianBlur(img, (0, 0), sigma)
+    G_large = cv.GaussianBlur(img, (0, 0), sigma_large)
     S = (1+p) * G_small - p * G_large
 
     return S
 
+
 def softThreshold(SI, epsilon, phi):
     T = np.zeros(SI.shape)
     SI_bright = SI >= epsilon
     SI_dark = SI < epsilon
     T[SI_bright] = 1.0
-    T[SI_dark] = 1.0 + np.tanh( phi * (SI[SI_dark] - epsilon))
+    T[SI_dark] = 1.0 + np.tanh(phi * (SI[SI_dark] - epsilon))
 
     return T
 
+
 def xdog(img, sigma, k_sigma, p, epsilon, phi):
     S = sharpImage(img, sigma, k_sigma, p)
     SI = np.multiply(img, S)
     T = softThreshold(SI, epsilon, phi)
 
     return T
 
-def line_process(filename):
+
+def xdog_process(filename):
     img = cv.imread(filename)
     img = cv.cvtColor(img, cv.COLOR_RGB2GRAY)
     img = img / 255.0
     sigma = np.random.choice([0.3, 0.4, 0.5])
-    img = xdog(img, sigma, 4.5, 19,0.01, 10^9)
+    img = xdog(img, sigma, 4.5, 19, 0.01, 10^9)
+    img[img < 0.9] = 0.0
     img = img * 255
     img = img.reshape(img.shape[0], img.shape[1], 1)
-    img = np.tile(img, (1,1,3))
+    img = np.tile(img, (1, 1, 3))
 
     return img