init

app.py

@@ -0,0 +1,88 @@
+#!/usr/bin/env python
+
+from __future__ import annotations
+import argparse
+import functools
+import os
+import pathlib
+import sys
+from typing import Callable
+
+
+import gradio as gr
+import huggingface_hub
+import numpy as np
+import PIL.Image
+
+from io import BytesIO
+
+
+ORIGINAL_REPO_URL = 'https://github.com/taki0112/UGATIT'
+TITLE = 'taki0112/UGATIT'
+DESCRIPTION = f"""This is a demo for {ORIGINAL_REPO_URL}.
+
+"""
+ARTICLE = """
+
+"""
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--device', type=str, default='cpu')
+    parser.add_argument('--theme', type=str)
+    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')
+    parser.add_argument('--allow-screenshot', action='store_true')
+    return parser.parse_args()
+
+
+
+def run(
+    image
+) -> tuple[PIL.Image.Image]:
+    
+
+    return PIL.Image.open(image.name)
+
+
+def main():
+    gr.close_all()
+
+    args = parse_args()
+
+    func = functools.partial(run)
+    func = functools.update_wrapper(func, run)
+
+    
+    gr.Interface(
+        func,
+        [
+            gr.inputs.Image(type='file', label='Input Image'),
+        ],
+        [
+            gr.outputs.Image(
+                type='pil',
+                label='Result'),
+        ],
+        #examples=examples,
+        theme=args.theme,
+        title=TITLE,
+        description=DESCRIPTION,
+        article=ARTICLE,
+        allow_screenshot=args.allow_screenshot,
+        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()

packages.txt

@@ -0,0 +1,2 @@
+
+

requirements.txt

@@ -0,0 +1,4 @@
+opencv-python-headless==4.5.5.62
+Pillow==9.0.1
+scipy==1.7.3
+tensorflow-gpu==1.14.0

ugatit/UGATIT.py

@@ -0,0 +1,665 @@
+from ugatit.ops import *
+from utils import *
+from glob import glob
+import time
+from tensorflow.contrib.data import prefetch_to_device, shuffle_and_repeat, map_and_batch
+import numpy as np
+
+class UGATIT(object) :
+    def __init__(self, sess, args):
+        self.light = args.light
+
+        if self.light :
+            self.model_name = 'UGATIT_light'
+        else :
+            self.model_name = 'UGATIT'
+
+        self.sess = sess
+        self.phase = args.phase
+        self.checkpoint_dir = args.checkpoint_dir
+        self.result_dir = args.result_dir
+        self.log_dir = args.log_dir
+        self.dataset_name = args.dataset
+        self.augment_flag = args.augment_flag
+
+        self.epoch = args.epoch
+        self.iteration = args.iteration
+        self.decay_flag = args.decay_flag
+        self.decay_epoch = args.decay_epoch
+
+        self.gan_type = args.gan_type
+
+        self.batch_size = args.batch_size
+        self.print_freq = args.print_freq
+        self.save_freq = args.save_freq
+
+        self.init_lr = args.lr
+        self.ch = args.ch
+
+        """ Weight """
+        self.adv_weight = args.adv_weight
+        self.cycle_weight = args.cycle_weight
+        self.identity_weight = args.identity_weight
+        self.cam_weight = args.cam_weight
+        self.ld = args.GP_ld
+        self.smoothing = args.smoothing
+
+        """ Generator """
+        self.n_res = args.n_res
+
+        """ Discriminator """
+        self.n_dis = args.n_dis
+        self.n_critic = args.n_critic
+        self.sn = args.sn
+
+        self.img_size = args.img_size
+        self.img_ch = args.img_ch
+
+
+        self.sample_dir = os.path.join(args.sample_dir, self.model_dir)
+        check_folder(self.sample_dir)
+
+        # self.trainA, self.trainB = prepare_data(dataset_name=self.dataset_name, size=self.img_size
+        self.trainA_dataset = glob('./dataset/{}/*.*'.format(self.dataset_name + '/trainA'))
+        self.trainB_dataset = glob('./dataset/{}/*.*'.format(self.dataset_name + '/trainB'))
+        self.dataset_num = max(len(self.trainA_dataset), len(self.trainB_dataset))
+
+        print()
+
+        print("##### Information #####")
+        print("# light : ", self.light)
+        print("# gan type : ", self.gan_type)
+        print("# dataset : ", self.dataset_name)
+        print("# max dataset number : ", self.dataset_num)
+        print("# batch_size : ", self.batch_size)
+        print("# epoch : ", self.epoch)
+        print("# iteration per epoch : ", self.iteration)
+        print("# smoothing : ", self.smoothing)
+
+        print()
+
+        print("##### Generator #####")
+        print("# residual blocks : ", self.n_res)
+
+        print()
+
+        print("##### Discriminator #####")
+        print("# discriminator layer : ", self.n_dis)
+        print("# the number of critic : ", self.n_critic)
+        print("# spectral normalization : ", self.sn)
+
+        print()
+
+        print("##### Weight #####")
+        print("# adv_weight : ", self.adv_weight)
+        print("# cycle_weight : ", self.cycle_weight)
+        print("# identity_weight : ", self.identity_weight)
+        print("# cam_weight : ", self.cam_weight)
+
+    ##################################################################################
+    # Generator
+    ##################################################################################
+
+    def generator(self, x_init, reuse=False, scope="generator"):
+        channel = self.ch
+        with tf.variable_scope(scope, reuse=reuse) :
+            x = conv(x_init, channel, kernel=7, stride=1, pad=3, pad_type='reflect', scope='conv')
+            x = instance_norm(x, scope='ins_norm')
+            x = relu(x)
+
+            # Down-Sampling
+            for i in range(2) :
+                x = conv(x, channel*2, kernel=3, stride=2, pad=1, pad_type='reflect', scope='conv_'+str(i))
+                x = instance_norm(x, scope='ins_norm_'+str(i))
+                x = relu(x)
+
+                channel = channel * 2
+
+            # Down-Sampling Bottleneck
+            for i in range(self.n_res):
+                x = resblock(x, channel, scope='resblock_' + str(i))
+
+
+            # Class Activation Map
+            cam_x = global_avg_pooling(x)
+            cam_gap_logit, cam_x_weight = fully_connected_with_w(cam_x, scope='CAM_logit')
+            x_gap = tf.multiply(x, cam_x_weight)
+
+            cam_x = global_max_pooling(x)
+            cam_gmp_logit, cam_x_weight = fully_connected_with_w(cam_x, reuse=True, scope='CAM_logit')
+            x_gmp = tf.multiply(x, cam_x_weight)
+
+
+            cam_logit = tf.concat([cam_gap_logit, cam_gmp_logit], axis=-1)
+            x = tf.concat([x_gap, x_gmp], axis=-1)
+
+            x = conv(x, channel, kernel=1, stride=1, scope='conv_1x1')
+            x = relu(x)
+
+            heatmap = tf.squeeze(tf.reduce_sum(x, axis=-1))
+
+            # Gamma, Beta block
+            gamma, beta = self.MLP(x, reuse=reuse)
+
+            # Up-Sampling Bottleneck
+            for i in range(self.n_res):
+                x = adaptive_ins_layer_resblock(x, channel, gamma, beta, smoothing=self.smoothing, scope='adaptive_resblock' + str(i))
+
+            # Up-Sampling
+            for i in range(2) :
+                x = up_sample(x, scale_factor=2)
+                x = conv(x, channel//2, kernel=3, stride=1, pad=1, pad_type='reflect', scope='up_conv_'+str(i))
+                x = layer_instance_norm(x, scope='layer_ins_norm_'+str(i))
+                x = relu(x)
+
+                channel = channel // 2
+
+
+            x = conv(x, channels=3, kernel=7, stride=1, pad=3, pad_type='reflect', scope='G_logit')
+            x = tanh(x)
+
+            return x, cam_logit, heatmap
+
+    def MLP(self, x, use_bias=True, reuse=False, scope='MLP'):
+        channel = self.ch * self.n_res
+
+        if self.light :
+            x = global_avg_pooling(x)
+
+        with tf.variable_scope(scope, reuse=reuse):
+            for i in range(2) :
+                x = fully_connected(x, channel, use_bias, scope='linear_' + str(i))
+                x = relu(x)
+
+
+            gamma = fully_connected(x, channel, use_bias, scope='gamma')
+            beta = fully_connected(x, channel, use_bias, scope='beta')
+
+            gamma = tf.reshape(gamma, shape=[self.batch_size, 1, 1, channel])
+            beta = tf.reshape(beta, shape=[self.batch_size, 1, 1, channel])
+
+            return gamma, beta
+
+    ##################################################################################
+    # Discriminator
+    ##################################################################################
+
+    def discriminator(self, x_init, reuse=False, scope="discriminator"):
+        D_logit = []
+        D_CAM_logit = []
+        with tf.variable_scope(scope, reuse=reuse) :
+            local_x, local_cam, local_heatmap = self.discriminator_local(x_init, reuse=reuse, scope='local')
+            global_x, global_cam, global_heatmap = self.discriminator_global(x_init, reuse=reuse, scope='global')
+
+            D_logit.extend([local_x, global_x])
+            D_CAM_logit.extend([local_cam, global_cam])
+
+            return D_logit, D_CAM_logit, local_heatmap, global_heatmap
+
+    def discriminator_global(self, x_init, reuse=False, scope='discriminator_global'):
+        with tf.variable_scope(scope, reuse=reuse):
+            channel = self.ch
+            x = conv(x_init, channel, kernel=4, stride=2, pad=1, pad_type='reflect', sn=self.sn, scope='conv_0')
+            x = lrelu(x, 0.2)
+
+            for i in range(1, self.n_dis - 1):
+                x = conv(x, channel * 2, kernel=4, stride=2, pad=1, pad_type='reflect', sn=self.sn, scope='conv_' + str(i))
+                x = lrelu(x, 0.2)
+
+                channel = channel * 2
+
+            x = conv(x, channel * 2, kernel=4, stride=1, pad=1, pad_type='reflect', sn=self.sn, scope='conv_last')
+            x = lrelu(x, 0.2)
+
+            channel = channel * 2
+
+            cam_x = global_avg_pooling(x)
+            cam_gap_logit, cam_x_weight = fully_connected_with_w(cam_x, sn=self.sn, scope='CAM_logit')
+            x_gap = tf.multiply(x, cam_x_weight)
+
+            cam_x = global_max_pooling(x)
+            cam_gmp_logit, cam_x_weight = fully_connected_with_w(cam_x, sn=self.sn, reuse=True, scope='CAM_logit')
+            x_gmp = tf.multiply(x, cam_x_weight)
+
+            cam_logit = tf.concat([cam_gap_logit, cam_gmp_logit], axis=-1)
+            x = tf.concat([x_gap, x_gmp], axis=-1)
+
+            x = conv(x, channel, kernel=1, stride=1, scope='conv_1x1')
+            x = lrelu(x, 0.2)
+
+            heatmap = tf.squeeze(tf.reduce_sum(x, axis=-1))
+
+
+            x = conv(x, channels=1, kernel=4, stride=1, pad=1, pad_type='reflect', sn=self.sn, scope='D_logit')
+
+            return x, cam_logit, heatmap
+
+    def discriminator_local(self, x_init, reuse=False, scope='discriminator_local'):
+        with tf.variable_scope(scope, reuse=reuse) :
+            channel = self.ch
+            x = conv(x_init, channel, kernel=4, stride=2, pad=1, pad_type='reflect', sn=self.sn, scope='conv_0')
+            x = lrelu(x, 0.2)
+
+            for i in range(1, self.n_dis - 2 - 1):
+                x = conv(x, channel * 2, kernel=4, stride=2, pad=1, pad_type='reflect', sn=self.sn, scope='conv_' + str(i))
+                x = lrelu(x, 0.2)
+
+                channel = channel * 2
+
+            x = conv(x, channel * 2, kernel=4, stride=1, pad=1, pad_type='reflect', sn=self.sn, scope='conv_last')
+            x = lrelu(x, 0.2)
+
+            channel = channel * 2
+
+            cam_x = global_avg_pooling(x)
+            cam_gap_logit, cam_x_weight = fully_connected_with_w(cam_x, sn=self.sn, scope='CAM_logit')
+            x_gap = tf.multiply(x, cam_x_weight)
+
+            cam_x = global_max_pooling(x)
+            cam_gmp_logit, cam_x_weight = fully_connected_with_w(cam_x, sn=self.sn, reuse=True, scope='CAM_logit')
+            x_gmp = tf.multiply(x, cam_x_weight)
+
+            cam_logit = tf.concat([cam_gap_logit, cam_gmp_logit], axis=-1)
+            x = tf.concat([x_gap, x_gmp], axis=-1)
+
+            x = conv(x, channel, kernel=1, stride=1, scope='conv_1x1')
+            x = lrelu(x, 0.2)
+
+            heatmap = tf.squeeze(tf.reduce_sum(x, axis=-1))
+
+            x = conv(x, channels=1, kernel=4, stride=1, pad=1, pad_type='reflect', sn=self.sn, scope='D_logit')
+
+            return x, cam_logit, heatmap
+
+    ##################################################################################
+    # Model
+    ##################################################################################
+
+    def generate_a2b(self, x_A, reuse=False):
+        out, cam, _ = self.generator(x_A, reuse=reuse, scope="generator_B")
+
+        return out, cam
+
+    def generate_b2a(self, x_B, reuse=False):
+        out, cam, _ = self.generator(x_B, reuse=reuse, scope="generator_A")
+
+        return out, cam
+
+    def discriminate_real(self, x_A, x_B):
+        real_A_logit, real_A_cam_logit, _, _ = self.discriminator(x_A, scope="discriminator_A")
+        real_B_logit, real_B_cam_logit, _, _ = self.discriminator(x_B, scope="discriminator_B")
+
+        return real_A_logit, real_A_cam_logit, real_B_logit, real_B_cam_logit
+
+    def discriminate_fake(self, x_ba, x_ab):
+        fake_A_logit, fake_A_cam_logit, _, _ = self.discriminator(x_ba, reuse=True, scope="discriminator_A")
+        fake_B_logit, fake_B_cam_logit, _, _ = self.discriminator(x_ab, reuse=True, scope="discriminator_B")
+
+        return fake_A_logit, fake_A_cam_logit, fake_B_logit, fake_B_cam_logit
+
+    def gradient_panalty(self, real, fake, scope="discriminator_A"):
+        if self.gan_type.__contains__('dragan'):
+            eps = tf.random_uniform(shape=tf.shape(real), minval=0., maxval=1.)
+            _, x_var = tf.nn.moments(real, axes=[0, 1, 2, 3])
+            x_std = tf.sqrt(x_var)  # magnitude of noise decides the size of local region
+
+            fake = real + 0.5 * x_std * eps
+
+        alpha = tf.random_uniform(shape=[self.batch_size, 1, 1, 1], minval=0., maxval=1.)
+        interpolated = real + alpha * (fake - real)
+
+        logit, cam_logit, _, _ = self.discriminator(interpolated, reuse=True, scope=scope)
+
+
+        GP = []
+        cam_GP = []
+
+        for i in range(2) :
+            grad = tf.gradients(logit[i], interpolated)[0] # gradient of D(interpolated)
+            grad_norm = tf.norm(flatten(grad), axis=1) # l2 norm
+
+            # WGAN - LP
+            if self.gan_type == 'wgan-lp' :
+                GP.append(self.ld * tf.reduce_mean(tf.square(tf.maximum(0.0, grad_norm - 1.))))
+
+            elif self.gan_type == 'wgan-gp' or self.gan_type == 'dragan':
+                GP.append(self.ld * tf.reduce_mean(tf.square(grad_norm - 1.)))
+
+        for i in range(2) :
+            grad = tf.gradients(cam_logit[i], interpolated)[0] # gradient of D(interpolated)
+            grad_norm = tf.norm(flatten(grad), axis=1) # l2 norm
+
+            # WGAN - LP
+            if self.gan_type == 'wgan-lp' :
+                cam_GP.append(self.ld * tf.reduce_mean(tf.square(tf.maximum(0.0, grad_norm - 1.))))
+
+            elif self.gan_type == 'wgan-gp' or self.gan_type == 'dragan':
+                cam_GP.append(self.ld * tf.reduce_mean(tf.square(grad_norm - 1.)))
+
+
+        return sum(GP), sum(cam_GP)
+
+    def build_model(self):
+        if self.phase == 'train' :
+            self.lr = tf.placeholder(tf.float32, name='learning_rate')
+
+
+            """ Input Image"""
+            Image_Data_Class = ImageData(self.img_size, self.img_ch, self.augment_flag)
+
+            trainA = tf.data.Dataset.from_tensor_slices(self.trainA_dataset)
+            trainB = tf.data.Dataset.from_tensor_slices(self.trainB_dataset)
+
+
+            gpu_device = '/gpu:0'
+            trainA = trainA.apply(shuffle_and_repeat(self.dataset_num)).apply(map_and_batch(Image_Data_Class.image_processing, self.batch_size, num_parallel_batches=16, drop_remainder=True)).apply(prefetch_to_device(gpu_device, None))
+            trainB = trainB.apply(shuffle_and_repeat(self.dataset_num)).apply(map_and_batch(Image_Data_Class.image_processing, self.batch_size, num_parallel_batches=16, drop_remainder=True)).apply(prefetch_to_device(gpu_device, None))
+
+
+            trainA_iterator = trainA.make_one_shot_iterator()
+            trainB_iterator = trainB.make_one_shot_iterator()
+
+            self.domain_A = trainA_iterator.get_next()
+            self.domain_B = trainB_iterator.get_next()
+
+            """ Define Generator, Discriminator """
+            x_ab, cam_ab = self.generate_a2b(self.domain_A) # real a
+            x_ba, cam_ba = self.generate_b2a(self.domain_B) # real b
+
+            x_aba, _ = self.generate_b2a(x_ab, reuse=True) # real b
+            x_bab, _ = self.generate_a2b(x_ba, reuse=True) # real a
+
+            x_aa, cam_aa = self.generate_b2a(self.domain_A, reuse=True) # fake b
+            x_bb, cam_bb = self.generate_a2b(self.domain_B, reuse=True) # fake a
+
+            real_A_logit, real_A_cam_logit, real_B_logit, real_B_cam_logit = self.discriminate_real(self.domain_A, self.domain_B)
+            fake_A_logit, fake_A_cam_logit, fake_B_logit, fake_B_cam_logit = self.discriminate_fake(x_ba, x_ab)
+
+
+            """ Define Loss """
+            if self.gan_type.__contains__('wgan') or self.gan_type == 'dragan' :
+                GP_A, GP_CAM_A = self.gradient_panalty(real=self.domain_A, fake=x_ba, scope="discriminator_A")
+                GP_B, GP_CAM_B = self.gradient_panalty(real=self.domain_B, fake=x_ab, scope="discriminator_B")
+            else :
+                GP_A, GP_CAM_A  = 0, 0
+                GP_B, GP_CAM_B = 0, 0
+
+            G_ad_loss_A = (generator_loss(self.gan_type, fake_A_logit) + generator_loss(self.gan_type, fake_A_cam_logit))
+            G_ad_loss_B = (generator_loss(self.gan_type, fake_B_logit) + generator_loss(self.gan_type, fake_B_cam_logit))
+
+            D_ad_loss_A = (discriminator_loss(self.gan_type, real_A_logit, fake_A_logit) + discriminator_loss(self.gan_type, real_A_cam_logit, fake_A_cam_logit) + GP_A + GP_CAM_A)
+            D_ad_loss_B = (discriminator_loss(self.gan_type, real_B_logit, fake_B_logit) + discriminator_loss(self.gan_type, real_B_cam_logit, fake_B_cam_logit) + GP_B + GP_CAM_B)
+
+            reconstruction_A = L1_loss(x_aba, self.domain_A) # reconstruction
+            reconstruction_B = L1_loss(x_bab, self.domain_B) # reconstruction
+
+            identity_A = L1_loss(x_aa, self.domain_A)
+            identity_B = L1_loss(x_bb, self.domain_B)
+
+            cam_A = cam_loss(source=cam_ba, non_source=cam_aa)
+            cam_B = cam_loss(source=cam_ab, non_source=cam_bb)
+
+            Generator_A_gan = self.adv_weight * G_ad_loss_A
+            Generator_A_cycle = self.cycle_weight * reconstruction_B
+            Generator_A_identity = self.identity_weight * identity_A
+            Generator_A_cam = self.cam_weight * cam_A
+
+
+            Generator_B_gan = self.adv_weight * G_ad_loss_B
+            Generator_B_cycle = self.cycle_weight * reconstruction_A
+            Generator_B_identity = self.identity_weight * identity_B
+            Generator_B_cam = self.cam_weight * cam_B
+
+
+            Generator_A_loss = Generator_A_gan + Generator_A_cycle + Generator_A_identity + Generator_A_cam
+            Generator_B_loss = Generator_B_gan + Generator_B_cycle + Generator_B_identity + Generator_B_cam
+
+
+            Discriminator_A_loss = self.adv_weight * D_ad_loss_A
+            Discriminator_B_loss = self.adv_weight * D_ad_loss_B
+
+            self.Generator_loss = Generator_A_loss + Generator_B_loss + regularization_loss('generator')
+            self.Discriminator_loss = Discriminator_A_loss + Discriminator_B_loss + regularization_loss('discriminator')
+
+
+            """ Result Image """
+            self.fake_A = x_ba
+            self.fake_B = x_ab
+
+            self.real_A = self.domain_A
+            self.real_B = self.domain_B
+
+
+            """ Training """
+            t_vars = tf.trainable_variables()
+            G_vars = [var for var in t_vars if 'generator' in var.name]
+            D_vars = [var for var in t_vars if 'discriminator' in var.name]
+
+            self.G_optim = tf.train.AdamOptimizer(self.lr, beta1=0.5, beta2=0.999).minimize(self.Generator_loss, var_list=G_vars)
+            self.D_optim = tf.train.AdamOptimizer(self.lr, beta1=0.5, beta2=0.999).minimize(self.Discriminator_loss, var_list=D_vars)
+
+
+            """" Summary """
+            self.all_G_loss = tf.summary.scalar("Generator_loss", self.Generator_loss)
+            self.all_D_loss = tf.summary.scalar("Discriminator_loss", self.Discriminator_loss)
+
+            self.G_A_loss = tf.summary.scalar("G_A_loss", Generator_A_loss)
+            self.G_A_gan = tf.summary.scalar("G_A_gan", Generator_A_gan)
+            self.G_A_cycle = tf.summary.scalar("G_A_cycle", Generator_A_cycle)
+            self.G_A_identity = tf.summary.scalar("G_A_identity", Generator_A_identity)
+            self.G_A_cam = tf.summary.scalar("G_A_cam", Generator_A_cam)
+
+            self.G_B_loss = tf.summary.scalar("G_B_loss", Generator_B_loss)
+            self.G_B_gan = tf.summary.scalar("G_B_gan", Generator_B_gan)
+            self.G_B_cycle = tf.summary.scalar("G_B_cycle", Generator_B_cycle)
+            self.G_B_identity = tf.summary.scalar("G_B_identity", Generator_B_identity)
+            self.G_B_cam = tf.summary.scalar("G_B_cam", Generator_B_cam)
+
+            self.D_A_loss = tf.summary.scalar("D_A_loss", Discriminator_A_loss)
+            self.D_B_loss = tf.summary.scalar("D_B_loss", Discriminator_B_loss)
+
+            self.rho_var = []
+            for var in tf.trainable_variables():
+                if 'rho' in var.name:
+                    self.rho_var.append(tf.summary.histogram(var.name, var))
+                    self.rho_var.append(tf.summary.scalar(var.name + "_min", tf.reduce_min(var)))
+                    self.rho_var.append(tf.summary.scalar(var.name + "_max", tf.reduce_max(var)))
+                    self.rho_var.append(tf.summary.scalar(var.name + "_mean", tf.reduce_mean(var)))
+
+            g_summary_list = [self.G_A_loss, self.G_A_gan, self.G_A_cycle, self.G_A_identity, self.G_A_cam,
+                              self.G_B_loss, self.G_B_gan, self.G_B_cycle, self.G_B_identity, self.G_B_cam,
+                              self.all_G_loss]
+
+            g_summary_list.extend(self.rho_var)
+            d_summary_list = [self.D_A_loss, self.D_B_loss, self.all_D_loss]
+
+            self.G_loss = tf.summary.merge(g_summary_list)
+            self.D_loss = tf.summary.merge(d_summary_list)
+
+        else :
+            """ Test """
+            self.test_domain_A = tf.placeholder(tf.float32, [1, self.img_size, self.img_size, self.img_ch], name='test_domain_A')
+            self.test_domain_B = tf.placeholder(tf.float32, [1, self.img_size, self.img_size, self.img_ch], name='test_domain_B')
+
+
+            self.test_fake_B, _ = self.generate_a2b(self.test_domain_A)
+            self.test_fake_A, _ = self.generate_b2a(self.test_domain_B)
+
+
+    def train(self):
+        # initialize all variables
+        tf.global_variables_initializer().run()
+
+        # saver to save model
+        self.saver = tf.train.Saver()
+
+        # summary writer
+        self.writer = tf.summary.FileWriter(self.log_dir + '/' + self.model_dir, self.sess.graph)
+
+
+        # restore check-point if it exits
+        could_load, checkpoint_counter = self.load(self.checkpoint_dir)
+        if could_load:
+            start_epoch = (int)(checkpoint_counter / self.iteration)
+            start_batch_id = checkpoint_counter - start_epoch * self.iteration
+            counter = checkpoint_counter
+            print(" [*] Load SUCCESS")
+        else:
+            start_epoch = 0
+            start_batch_id = 0
+            counter = 1
+            print(" [!] Load failed...")
+
+        # loop for epoch
+        start_time = time.time()
+        past_g_loss = -1.
+        lr = self.init_lr
+        for epoch in range(start_epoch, self.epoch):
+            # lr = self.init_lr if epoch < self.decay_epoch else self.init_lr * (self.epoch - epoch) / (self.epoch - self.decay_epoch)
+            if self.decay_flag :
+                #lr = self.init_lr * pow(0.5, epoch // self.decay_epoch)
+                lr = self.init_lr if epoch < self.decay_epoch else self.init_lr * (self.epoch - epoch) / (self.epoch - self.decay_epoch)
+            for idx in range(start_batch_id, self.iteration):
+                train_feed_dict = {
+                    self.lr : lr
+                }
+
+                # Update D
+                _, d_loss, summary_str = self.sess.run([self.D_optim,
+                                                        self.Discriminator_loss, self.D_loss], feed_dict = train_feed_dict)
+                self.writer.add_summary(summary_str, counter)
+
+                # Update G
+                g_loss = None
+                if (counter - 1) % self.n_critic == 0 :
+                    batch_A_images, batch_B_images, fake_A, fake_B, _, g_loss, summary_str = self.sess.run([self.real_A, self.real_B,
+                                                                                                            self.fake_A, self.fake_B,
+                                                                                                            self.G_optim,
+                                                                                                            self.Generator_loss, self.G_loss], feed_dict = train_feed_dict)
+                    self.writer.add_summary(summary_str, counter)
+                    past_g_loss = g_loss
+
+                # display training status
+                counter += 1
+                if g_loss == None :
+                    g_loss = past_g_loss
+                print("Epoch: [%2d] [%5d/%5d] time: %4.4f d_loss: %.8f, g_loss: %.8f" % (epoch, idx, self.iteration, time.time() - start_time, d_loss, g_loss))
+
+                if np.mod(idx+1, self.print_freq) == 0 :
+                    save_images(batch_A_images, [self.batch_size, 1],
+                                './{}/real_A_{:03d}_{:05d}.png'.format(self.sample_dir, epoch, idx+1))
+                    # save_images(batch_B_images, [self.batch_size, 1],
+                    #             './{}/real_B_{:03d}_{:05d}.png'.format(self.sample_dir, epoch, idx+1))
+
+                    # save_images(fake_A, [self.batch_size, 1],
+                    #             './{}/fake_A_{:03d}_{:05d}.png'.format(self.sample_dir, epoch, idx+1))
+                    save_images(fake_B, [self.batch_size, 1],
+                                './{}/fake_B_{:03d}_{:05d}.png'.format(self.sample_dir, epoch, idx+1))
+
+                if np.mod(idx + 1, self.save_freq) == 0:
+                    self.save(self.checkpoint_dir, counter)
+
+
+
+            # After an epoch, start_batch_id is set to zero
+            # non-zero value is only for the first epoch after loading pre-trained model
+            start_batch_id = 0
+
+            # save model for final step
+            self.save(self.checkpoint_dir, counter)
+
+    @property
+    def model_dir(self):
+        n_res = str(self.n_res) + 'resblock'
+        n_dis = str(self.n_dis) + 'dis'
+
+        if self.smoothing :
+            smoothing = '_smoothing'
+        else :
+            smoothing = ''
+
+        if self.sn :
+            sn = '_sn'
+        else :
+            sn = ''
+
+        return "{}_{}_{}_{}_{}_{}_{}_{}_{}_{}{}{}".format(self.model_name, self.dataset_name,
+                                                         self.gan_type, n_res, n_dis,
+                                                         self.n_critic,
+                                                         self.adv_weight, self.cycle_weight, self.identity_weight, self.cam_weight, sn, smoothing)
+
+    def save(self, checkpoint_dir, step):
+        checkpoint_dir = os.path.join(checkpoint_dir, self.model_dir)
+
+        if not os.path.exists(checkpoint_dir):
+            os.makedirs(checkpoint_dir)
+
+        self.saver.save(self.sess, os.path.join(checkpoint_dir, self.model_name + '.model'), global_step=step)
+
+    def load(self, checkpoint_dir):
+        print(" [*] Reading checkpoints...")
+        checkpoint_dir = os.path.join(checkpoint_dir, self.model_dir)
+
+        ckpt = tf.train.get_checkpoint_state(checkpoint_dir)
+        if ckpt and ckpt.model_checkpoint_path:
+            ckpt_name = os.path.basename(ckpt.model_checkpoint_path)
+            self.saver.restore(self.sess, os.path.join(checkpoint_dir, ckpt_name))
+            counter = int(ckpt_name.split('-')[-1])
+            print(" [*] Success to read {}".format(ckpt_name))
+            return True, counter
+        else:
+            print(" [*] Failed to find a checkpoint")
+            return False, 0
+
+    def test(self):
+        tf.global_variables_initializer().run()
+        test_A_files = glob('./dataset/{}/*.*'.format(self.dataset_name + '/testA'))
+        test_B_files = glob('./dataset/{}/*.*'.format(self.dataset_name + '/testB'))
+
+        self.saver = tf.train.Saver()
+        could_load, checkpoint_counter = self.load(self.checkpoint_dir)
+        self.result_dir = os.path.join(self.result_dir, self.model_dir)
+        check_folder(self.result_dir)
+
+        if could_load :
+            print(" [*] Load SUCCESS")
+        else :
+            print(" [!] Load failed...")
+
+        # write html for visual comparison
+        index_path = os.path.join(self.result_dir, 'index.html')
+        index = open(index_path, 'w')
+        index.write("<html><body><table><tr>")
+        index.write("<th>name</th><th>input</th><th>output</th></tr>")
+
+        for sample_file  in test_A_files : # A -> B
+            print('Processing A image: ' + sample_file)
+            sample_image = np.asarray(load_test_data(sample_file, size=self.img_size))
+            image_path = os.path.join(self.result_dir,'{0}'.format(os.path.basename(sample_file)))
+
+            fake_img = self.sess.run(self.test_fake_B, feed_dict = {self.test_domain_A : sample_image})
+            save_images(fake_img, [1, 1], image_path)
+
+            index.write("<td>%s</td>" % os.path.basename(image_path))
+
+            index.write("<td><img src='%s' width='%d' height='%d'></td>" % (sample_file if os.path.isabs(sample_file) else (
+                '../..' + os.path.sep + sample_file), self.img_size, self.img_size))
+            index.write("<td><img src='%s' width='%d' height='%d'></td>" % (image_path if os.path.isabs(image_path) else (
+                '../..' + os.path.sep + image_path), self.img_size, self.img_size))
+            index.write("</tr>")
+
+        for sample_file  in test_B_files : # B -> A
+            print('Processing B image: ' + sample_file)
+            sample_image = np.asarray(load_test_data(sample_file, size=self.img_size))
+            image_path = os.path.join(self.result_dir,'{0}'.format(os.path.basename(sample_file)))
+
+            fake_img = self.sess.run(self.test_fake_A, feed_dict = {self.test_domain_B : sample_image})
+
+            save_images(fake_img, [1, 1], image_path)
+            index.write("<td>%s</td>" % os.path.basename(image_path))
+            index.write("<td><img src='%s' width='%d' height='%d'></td>" % (sample_file if os.path.isabs(sample_file) else (
+                    '../..' + os.path.sep + sample_file), self.img_size, self.img_size))
+            index.write("<td><img src='%s' width='%d' height='%d'></td>" % (image_path if os.path.isabs(image_path) else (
+                    '../..' + os.path.sep + image_path), self.img_size, self.img_size))
+            index.write("</tr>")
+        index.close()

ugatit/main.py

@@ -0,0 +1,106 @@
+from ugatit.UGATIT import UGATIT
+import argparse
+from ugatit.utils import *
+
+"""parsing and configuration"""
+
+def parse_args():
+    desc = "Tensorflow implementation of U-GAT-IT"
+    parser = argparse.ArgumentParser(description=desc)
+    parser.add_argument('--phase', type=str, default='train', help='[train / test]')
+    parser.add_argument('--light', type=str2bool, default=False, help='[U-GAT-IT full version / U-GAT-IT light version]')
+    parser.add_argument('--dataset', type=str, default='selfie2anime', help='dataset_name')
+
+    parser.add_argument('--epoch', type=int, default=100, help='The number of epochs to run')
+    parser.add_argument('--iteration', type=int, default=10000, help='The number of training iterations')
+    parser.add_argument('--batch_size', type=int, default=1, help='The size of batch size')
+    parser.add_argument('--print_freq', type=int, default=1000, help='The number of image_print_freq')
+    parser.add_argument('--save_freq', type=int, default=1000, help='The number of ckpt_save_freq')
+    parser.add_argument('--decay_flag', type=str2bool, default=True, help='The decay_flag')
+    parser.add_argument('--decay_epoch', type=int, default=50, help='decay epoch')
+
+    parser.add_argument('--lr', type=float, default=0.0001, help='The learning rate')
+    parser.add_argument('--GP_ld', type=int, default=10, help='The gradient penalty lambda')
+    parser.add_argument('--adv_weight', type=int, default=1, help='Weight about GAN')
+    parser.add_argument('--cycle_weight', type=int, default=10, help='Weight about Cycle')
+    parser.add_argument('--identity_weight', type=int, default=10, help='Weight about Identity')
+    parser.add_argument('--cam_weight', type=int, default=1000, help='Weight about CAM')
+    parser.add_argument('--gan_type', type=str, default='lsgan', help='[gan / lsgan / wgan-gp / wgan-lp / dragan / hinge]')
+
+    parser.add_argument('--smoothing', type=str2bool, default=True, help='AdaLIN smoothing effect')
+
+    parser.add_argument('--ch', type=int, default=64, help='base channel number per layer')
+    parser.add_argument('--n_res', type=int, default=4, help='The number of resblock')
+    parser.add_argument('--n_dis', type=int, default=6, help='The number of discriminator layer')
+    parser.add_argument('--n_critic', type=int, default=1, help='The number of critic')
+    parser.add_argument('--sn', type=str2bool, default=True, help='using spectral norm')
+
+    parser.add_argument('--img_size', type=int, default=256, help='The size of image')
+    parser.add_argument('--img_ch', type=int, default=3, help='The size of image channel')
+    parser.add_argument('--augment_flag', type=str2bool, default=True, help='Image augmentation use or not')
+
+    parser.add_argument('--checkpoint_dir', type=str, default='checkpoint',
+                        help='Directory name to save the checkpoints')
+    parser.add_argument('--result_dir', type=str, default='results',
+                        help='Directory name to save the generated images')
+    parser.add_argument('--log_dir', type=str, default='logs',
+                        help='Directory name to save training logs')
+    parser.add_argument('--sample_dir', type=str, default='samples',
+                        help='Directory name to save the samples on training')
+
+    return check_args(parser.parse_args())
+
+"""checking arguments"""
+def check_args(args):
+    # --checkpoint_dir
+    check_folder(args.checkpoint_dir)
+
+    # --result_dir
+    check_folder(args.result_dir)
+
+    # --result_dir
+    check_folder(args.log_dir)
+
+    # --sample_dir
+    check_folder(args.sample_dir)
+
+    # --epoch
+    try:
+        assert args.epoch >= 1
+    except:
+        print('number of epochs must be larger than or equal to one')
+
+    # --batch_size
+    try:
+        assert args.batch_size >= 1
+    except:
+        print('batch size must be larger than or equal to one')
+    return args
+
+"""main"""
+def main():
+    # parse arguments
+    args = parse_args()
+    if args is None:
+      exit()
+
+    # open session
+    with tf.Session(config=tf.ConfigProto(allow_soft_placement=True)) as sess:
+        gan = UGATIT(sess, args)
+
+        # build graph
+        gan.build_model()
+
+        # show network architecture
+        show_all_variables()
+
+        if args.phase == 'train' :
+            gan.train()
+            print(" [*] Training finished!")
+
+        if args.phase == 'test' :
+            gan.test()
+            print(" [*] Test finished!")
+
+if __name__ == '__main__':
+    main()

ugatit/ops.py

@@ -0,0 +1,345 @@
+import tensorflow as tf
+import tensorflow.contrib as tf_contrib
+
+# Xavier : tf_contrib.layers.xavier_initializer()
+# He : tf_contrib.layers.variance_scaling_initializer()
+# Normal : tf.random_normal_initializer(mean=0.0, stddev=0.02)
+# l2_decay : tf_contrib.layers.l2_regularizer(0.0001)
+
+weight_init = tf.random_normal_initializer(mean=0.0, stddev=0.02)
+weight_regularizer = tf_contrib.layers.l2_regularizer(scale=0.0001)
+
+##################################################################################
+# Layer
+##################################################################################
+
+def conv(x, channels, kernel=4, stride=2, pad=0, pad_type='zero', use_bias=True, sn=False, scope='conv_0'):
+    with tf.variable_scope(scope):
+        if pad > 0 :
+            if (kernel - stride) % 2 == 0:
+                pad_top = pad
+                pad_bottom = pad
+                pad_left = pad
+                pad_right = pad
+
+            else:
+                pad_top = pad
+                pad_bottom = kernel - stride - pad_top
+                pad_left = pad
+                pad_right = kernel - stride - pad_left
+
+            if pad_type == 'zero':
+                x = tf.pad(x, [[0, 0], [pad_top, pad_bottom], [pad_left, pad_right], [0, 0]])
+            if pad_type == 'reflect':
+                x = tf.pad(x, [[0, 0], [pad_top, pad_bottom], [pad_left, pad_right], [0, 0]], mode='REFLECT')
+
+        if sn :
+            w = tf.get_variable("kernel", shape=[kernel, kernel, x.get_shape()[-1], channels], initializer=weight_init,
+                                regularizer=weight_regularizer)
+            x = tf.nn.conv2d(input=x, filter=spectral_norm(w),
+                             strides=[1, stride, stride, 1], padding='VALID')
+            if use_bias :
+                bias = tf.get_variable("bias", [channels], initializer=tf.constant_initializer(0.0))
+                x = tf.nn.bias_add(x, bias)
+
+        else :
+            x = tf.layers.conv2d(inputs=x, filters=channels,
+                                 kernel_size=kernel, kernel_initializer=weight_init,
+                                 kernel_regularizer=weight_regularizer,
+                                 strides=stride, use_bias=use_bias)
+
+
+        return x
+
+def fully_connected_with_w(x, use_bias=True, sn=False, reuse=False, scope='linear'):
+    with tf.variable_scope(scope, reuse=reuse):
+        x = flatten(x)
+        bias = 0.0
+        shape = x.get_shape().as_list()
+        channels = shape[-1]
+
+        w = tf.get_variable("kernel", [channels, 1], tf.float32,
+                            initializer=weight_init, regularizer=weight_regularizer)
+
+        if sn :
+            w = spectral_norm(w)
+
+        if use_bias :
+            bias = tf.get_variable("bias", [1],
+                                   initializer=tf.constant_initializer(0.0))
+
+            x = tf.matmul(x, w) + bias
+        else :
+            x = tf.matmul(x, w)
+
+        if use_bias :
+            weights = tf.gather(tf.transpose(tf.nn.bias_add(w, bias)), 0)
+        else :
+            weights = tf.gather(tf.transpose(w), 0)
+
+        return x, weights
+
+def fully_connected(x, units, use_bias=True, sn=False, scope='linear'):
+    with tf.variable_scope(scope):
+        x = flatten(x)
+        shape = x.get_shape().as_list()
+        channels = shape[-1]
+
+        if sn:
+            w = tf.get_variable("kernel", [channels, units], tf.float32,
+                                initializer=weight_init, regularizer=weight_regularizer)
+            if use_bias:
+                bias = tf.get_variable("bias", [units],
+                                       initializer=tf.constant_initializer(0.0))
+
+                x = tf.matmul(x, spectral_norm(w)) + bias
+            else:
+                x = tf.matmul(x, spectral_norm(w))
+
+        else :
+            x = tf.layers.dense(x, units=units, kernel_initializer=weight_init, kernel_regularizer=weight_regularizer, use_bias=use_bias)
+
+        return x
+
+def flatten(x) :
+    return tf.layers.flatten(x)
+
+##################################################################################
+# Residual-block
+##################################################################################
+
+def resblock(x_init, channels, use_bias=True, scope='resblock_0'):
+    with tf.variable_scope(scope):
+        with tf.variable_scope('res1'):
+            x = conv(x_init, channels, kernel=3, stride=1, pad=1, pad_type='reflect', use_bias=use_bias)
+            x = instance_norm(x)
+            x = relu(x)
+
+        with tf.variable_scope('res2'):
+            x = conv(x, channels, kernel=3, stride=1, pad=1, pad_type='reflect', use_bias=use_bias)
+            x = instance_norm(x)
+
+        return x + x_init
+
+def adaptive_ins_layer_resblock(x_init, channels, gamma, beta, use_bias=True, smoothing=True, scope='adaptive_resblock') :
+    with tf.variable_scope(scope):
+        with tf.variable_scope('res1'):
+            x = conv(x_init, channels, kernel=3, stride=1, pad=1, pad_type='reflect', use_bias=use_bias)
+            x = adaptive_instance_layer_norm(x, gamma, beta, smoothing)
+            x = relu(x)
+
+        with tf.variable_scope('res2'):
+            x = conv(x, channels, kernel=3, stride=1, pad=1, pad_type='reflect', use_bias=use_bias)
+            x = adaptive_instance_layer_norm(x, gamma, beta, smoothing)
+
+        return x + x_init
+
+
+##################################################################################
+# Sampling
+##################################################################################
+
+def up_sample(x, scale_factor=2):
+    _, h, w, _ = x.get_shape().as_list()
+    new_size = [h * scale_factor, w * scale_factor]
+    return tf.image.resize_nearest_neighbor(x, size=new_size)
+
+
+def global_avg_pooling(x):
+    gap = tf.reduce_mean(x, axis=[1, 2])
+    return gap
+
+def global_max_pooling(x):
+    gmp = tf.reduce_max(x, axis=[1, 2])
+    return gmp
+
+##################################################################################
+# Activation function
+##################################################################################
+
+def lrelu(x, alpha=0.01):
+    # pytorch alpha is 0.01
+    return tf.nn.leaky_relu(x, alpha)
+
+
+def relu(x):
+    return tf.nn.relu(x)
+
+
+def tanh(x):
+    return tf.tanh(x)
+
+def sigmoid(x) :
+    return tf.sigmoid(x)
+
+##################################################################################
+# Normalization function
+##################################################################################
+
+def adaptive_instance_layer_norm(x, gamma, beta, smoothing=True, scope='instance_layer_norm') :
+    with tf.variable_scope(scope):
+        ch = x.shape[-1]
+        eps = 1e-5
+
+        ins_mean, ins_sigma = tf.nn.moments(x, axes=[1, 2], keep_dims=True)
+        x_ins = (x - ins_mean) / (tf.sqrt(ins_sigma + eps))
+
+        ln_mean, ln_sigma = tf.nn.moments(x, axes=[1, 2, 3], keep_dims=True)
+        x_ln = (x - ln_mean) / (tf.sqrt(ln_sigma + eps))
+
+        rho = tf.get_variable("rho", [ch], initializer=tf.constant_initializer(1.0), constraint=lambda x: tf.clip_by_value(x, clip_value_min=0.0, clip_value_max=1.0))
+
+        if smoothing :
+            rho = tf.clip_by_value(rho - tf.constant(0.1), 0.0, 1.0)
+
+        x_hat = rho * x_ins + (1 - rho) * x_ln
+
+
+        x_hat = x_hat * gamma + beta
+
+        return x_hat
+
+def instance_norm(x, scope='instance_norm'):
+    return tf_contrib.layers.instance_norm(x,
+                                           epsilon=1e-05,
+                                           center=True, scale=True,
+                                           scope=scope)
+
+def layer_norm(x, scope='layer_norm') :
+    return tf_contrib.layers.layer_norm(x,
+                                        center=True, scale=True,
+                                        scope=scope)
+
+def layer_instance_norm(x, scope='layer_instance_norm') :
+    with tf.variable_scope(scope):
+        ch = x.shape[-1]
+        eps = 1e-5
+
+        ins_mean, ins_sigma = tf.nn.moments(x, axes=[1, 2], keep_dims=True)
+        x_ins = (x - ins_mean) / (tf.sqrt(ins_sigma + eps))
+
+        ln_mean, ln_sigma = tf.nn.moments(x, axes=[1, 2, 3], keep_dims=True)
+        x_ln = (x - ln_mean) / (tf.sqrt(ln_sigma + eps))
+
+        rho = tf.get_variable("rho", [ch], initializer=tf.constant_initializer(0.0), constraint=lambda x: tf.clip_by_value(x, clip_value_min=0.0, clip_value_max=1.0))
+
+        gamma = tf.get_variable("gamma", [ch], initializer=tf.constant_initializer(1.0))
+        beta = tf.get_variable("beta", [ch], initializer=tf.constant_initializer(0.0))
+
+        x_hat = rho * x_ins + (1 - rho) * x_ln
+
+        x_hat = x_hat * gamma + beta
+
+        return x_hat
+
+def spectral_norm(w, iteration=1):
+    w_shape = w.shape.as_list()
+    w = tf.reshape(w, [-1, w_shape[-1]])
+
+    u = tf.get_variable("u", [1, w_shape[-1]], initializer=tf.random_normal_initializer(), trainable=False)
+
+    u_hat = u
+    v_hat = None
+    for i in range(iteration):
+        """
+        power iteration
+        Usually iteration = 1 will be enough
+        """
+        v_ = tf.matmul(u_hat, tf.transpose(w))
+        v_hat = tf.nn.l2_normalize(v_)
+
+        u_ = tf.matmul(v_hat, w)
+        u_hat = tf.nn.l2_normalize(u_)
+
+    u_hat = tf.stop_gradient(u_hat)
+    v_hat = tf.stop_gradient(v_hat)
+
+    sigma = tf.matmul(tf.matmul(v_hat, w), tf.transpose(u_hat))
+
+    with tf.control_dependencies([u.assign(u_hat)]):
+        w_norm = w / sigma
+        w_norm = tf.reshape(w_norm, w_shape)
+
+
+    return w_norm
+
+##################################################################################
+# Loss function
+##################################################################################
+
+def L1_loss(x, y):
+    loss = tf.reduce_mean(tf.abs(x - y))
+
+    return loss
+
+def cam_loss(source, non_source) :
+
+    identity_loss = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(labels=tf.ones_like(source), logits=source))
+    non_identity_loss = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(labels=tf.zeros_like(non_source), logits=non_source))
+
+    loss = identity_loss + non_identity_loss
+
+    return loss
+
+def regularization_loss(scope_name) :
+    """
+    If you want to use "Regularization"
+    g_loss += regularization_loss('generator')
+    d_loss += regularization_loss('discriminator')
+    """
+    collection_regularization = tf.get_collection(tf.GraphKeys.REGULARIZATION_LOSSES)
+
+    loss = []
+    for item in collection_regularization :
+        if scope_name in item.name :
+            loss.append(item)
+
+    return tf.reduce_sum(loss)
+
+
+def discriminator_loss(loss_func, real, fake):
+    loss = []
+    real_loss = 0
+    fake_loss = 0
+
+    for i in range(2) :
+        if loss_func.__contains__('wgan') :
+            real_loss = -tf.reduce_mean(real[i])
+            fake_loss = tf.reduce_mean(fake[i])
+
+        if loss_func == 'lsgan' :
+            real_loss = tf.reduce_mean(tf.squared_difference(real[i], 1.0))
+            fake_loss = tf.reduce_mean(tf.square(fake[i]))
+
+        if loss_func == 'gan' or loss_func == 'dragan' :
+            real_loss = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(labels=tf.ones_like(real[i]), logits=real[i]))
+            fake_loss = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(labels=tf.zeros_like(fake[i]), logits=fake[i]))
+
+        if loss_func == 'hinge' :
+            real_loss = tf.reduce_mean(relu(1.0 - real[i]))
+            fake_loss = tf.reduce_mean(relu(1.0 + fake[i]))
+
+        loss.append(real_loss + fake_loss)
+
+    return sum(loss)
+
+def generator_loss(loss_func, fake):
+    loss = []
+    fake_loss = 0
+
+    for i in range(2) :
+        if loss_func.__contains__('wgan') :
+            fake_loss = -tf.reduce_mean(fake[i])
+
+        if loss_func == 'lsgan' :
+            fake_loss = tf.reduce_mean(tf.squared_difference(fake[i], 1.0))
+
+        if loss_func == 'gan' or loss_func == 'dragan' :
+            fake_loss = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(labels=tf.ones_like(fake[i]), logits=fake[i]))
+
+        if loss_func == 'hinge' :
+            fake_loss = -tf.reduce_mean(fake[i])
+
+        loss.append(fake_loss)
+
+    return sum(loss)

ugatit/utils.py

@@ -0,0 +1,80 @@
+import tensorflow as tf
+from tensorflow.contrib import slim
+import cv2
+import os, random
+import numpy as np
+
+class ImageData:
+
+    def __init__(self, load_size, channels, augment_flag):
+        self.load_size = load_size
+        self.channels = channels
+        self.augment_flag = augment_flag
+
+    def image_processing(self, filename):
+        x = tf.read_file(filename)
+        x_decode = tf.image.decode_jpeg(x, channels=self.channels)
+        img = tf.image.resize_images(x_decode, [self.load_size, self.load_size])
+        img = tf.cast(img, tf.float32) / 127.5 - 1
+
+        if self.augment_flag :
+            augment_size = self.load_size + (30 if self.load_size == 256 else 15)
+            p = random.random()
+            if p > 0.5:
+                img = augmentation(img, augment_size)
+
+        return img
+
+def load_test_data(image_path, size=256):
+    img = cv2.imread(image_path, flags=cv2.IMREAD_COLOR)
+    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+
+    img = cv2.resize(img, dsize=(size, size))
+
+    img = np.expand_dims(img, axis=0)
+    img = img/127.5 - 1
+
+    return img
+
+def augmentation(image, augment_size):
+    seed = random.randint(0, 2 ** 31 - 1)
+    ori_image_shape = tf.shape(image)
+    image = tf.image.random_flip_left_right(image, seed=seed)
+    image = tf.image.resize_images(image, [augment_size, augment_size])
+    image = tf.random_crop(image, ori_image_shape, seed=seed)
+    return image
+
+def save_images(images, size, image_path):
+    return imsave(inverse_transform(images), size, image_path)
+
+def inverse_transform(images):
+    return ((images+1.) / 2) * 255.0
+
+
+def imsave(images, size, path):
+    images = merge(images, size)
+    images = cv2.cvtColor(images.astype('uint8'), cv2.COLOR_RGB2BGR)
+
+    return cv2.imwrite(path, images)
+
+def merge(images, size):
+    h, w = images.shape[1], images.shape[2]
+    img = np.zeros((h * size[0], w * size[1], 3))
+    for idx, image in enumerate(images):
+        i = idx % size[1]
+        j = idx // size[1]
+        img[h*j:h*(j+1), w*i:w*(i+1), :] = image
+
+    return img
+
+def show_all_variables():
+    model_vars = tf.trainable_variables()
+    slim.model_analyzer.analyze_vars(model_vars, print_info=True)
+
+def check_folder(log_dir):
+    if not os.path.exists(log_dir):
+        os.makedirs(log_dir)
+    return log_dir
+
+def str2bool(x):
+    return x.lower() in ('true')

ugatit_test.py

@@ -0,0 +1,372 @@
+from ugatit.ops import *
+from ugatit.utils import *
+from glob import glob
+import time
+from tensorflow.contrib.data import prefetch_to_device, shuffle_and_repeat, map_and_batch
+import numpy as np
+from ugatit.utils import *
+
+class UgatitTest:
+
+    def __init__(self, sess):
+        self.light = False
+
+        if self.light:
+            self.model_name = 'UGATIT_light'
+        else:
+            self.model_name = 'UGATIT'
+
+        self.sess = sess
+        self.phase = 'test'
+        self.checkpoint_dir = '/home/hylee/cartoon/UGATIT/checkpoint'
+        self.result_dir = 'results'
+        self.log_dir = 'logs'
+        self.dataset_name = 'selfie2anime'
+        self.augment_flag = True
+
+        self.epoch = 100
+        self.iteration = 10000
+        self.decay_flag = True
+        self.decay_epoch = 50
+
+        self.gan_type = 'lsgan'
+
+        self.batch_size = 1
+        self.print_freq = 1000
+        self.save_freq = 1000
+
+        self.init_lr = 0.0001
+        self.ch = 64
+
+        """ Weight """
+        self.adv_weight = 1
+        self.cycle_weight = 10
+        self.identity_weight = 10
+        self.cam_weight = 1000
+        self.ld = 10
+        self.smoothing = True
+
+        """ Generator """
+        self.n_res = 4
+
+        """ Discriminator """
+        self.n_dis = 6
+        self.n_critic = 1
+        self.sn = True
+
+        self.img_size = 256
+        self.img_ch = 3
+
+        self.sample_dir = os.path.join('/home/hylee/cartoon/UGATIT/samples', self.model_dir)
+        check_folder(self.sample_dir)
+
+        # self.trainA, self.trainB = prepare_data(dataset_name=self.dataset_name, size=self.img_size
+        self.trainA_dataset = glob('./dataset/{}/*.*'.format(self.dataset_name + '/trainA'))
+        self.trainB_dataset = glob('./dataset/{}/*.*'.format(self.dataset_name + '/trainB'))
+        self.dataset_num = max(len(self.trainA_dataset), len(self.trainB_dataset))
+
+        print()
+
+        print("##### Information #####")
+        print("# light : ", self.light)
+        print("# gan type : ", self.gan_type)
+        print("# dataset : ", self.dataset_name)
+        print("# max dataset number : ", self.dataset_num)
+        print("# batch_size : ", self.batch_size)
+        print("# epoch : ", self.epoch)
+        print("# iteration per epoch : ", self.iteration)
+        print("# smoothing : ", self.smoothing)
+
+        print()
+
+        print("##### Generator #####")
+        print("# residual blocks : ", self.n_res)
+
+        print()
+
+        print("##### Discriminator #####")
+        print("# discriminator layer : ", self.n_dis)
+        print("# the number of critic : ", self.n_critic)
+        print("# spectral normalization : ", self.sn)
+
+        print()
+
+        print("##### Weight #####")
+        print("# adv_weight : ", self.adv_weight)
+        print("# cycle_weight : ", self.cycle_weight)
+        print("# identity_weight : ", self.identity_weight)
+        print("# cam_weight : ", self.cam_weight)
+
+    ##################################################################################
+    # Generator
+    ##################################################################################
+
+    def generator(self, x_init, reuse=False, scope="generator"):
+        channel = self.ch
+        with tf.variable_scope(scope, reuse=reuse) :
+            x = conv(x_init, channel, kernel=7, stride=1, pad=3, pad_type='reflect', scope='conv')
+            x = instance_norm(x, scope='ins_norm')
+            x = relu(x)
+
+            # Down-Sampling
+            for i in range(2) :
+                x = conv(x, channel*2, kernel=3, stride=2, pad=1, pad_type='reflect', scope='conv_'+str(i))
+                x = instance_norm(x, scope='ins_norm_'+str(i))
+                x = relu(x)
+
+                channel = channel * 2
+
+            # Down-Sampling Bottleneck
+            for i in range(self.n_res):
+                x = resblock(x, channel, scope='resblock_' + str(i))
+
+
+            # Class Activation Map
+            cam_x = global_avg_pooling(x)
+            cam_gap_logit, cam_x_weight = fully_connected_with_w(cam_x, scope='CAM_logit')
+            x_gap = tf.multiply(x, cam_x_weight)
+
+            cam_x = global_max_pooling(x)
+            cam_gmp_logit, cam_x_weight = fully_connected_with_w(cam_x, reuse=True, scope='CAM_logit')
+            x_gmp = tf.multiply(x, cam_x_weight)
+
+
+            cam_logit = tf.concat([cam_gap_logit, cam_gmp_logit], axis=-1)
+            x = tf.concat([x_gap, x_gmp], axis=-1)
+
+            x = conv(x, channel, kernel=1, stride=1, scope='conv_1x1')
+            x = relu(x)
+
+            heatmap = tf.squeeze(tf.reduce_sum(x, axis=-1))
+
+            # Gamma, Beta block
+            gamma, beta = self.MLP(x, reuse=reuse)
+
+            # Up-Sampling Bottleneck
+            for i in range(self.n_res):
+                x = adaptive_ins_layer_resblock(x, channel, gamma, beta, smoothing=self.smoothing, scope='adaptive_resblock' + str(i))
+
+            # Up-Sampling
+            for i in range(2) :
+                x = up_sample(x, scale_factor=2)
+                x = conv(x, channel//2, kernel=3, stride=1, pad=1, pad_type='reflect', scope='up_conv_'+str(i))
+                x = layer_instance_norm(x, scope='layer_ins_norm_'+str(i))
+                x = relu(x)
+
+                channel = channel // 2
+
+
+            x = conv(x, channels=3, kernel=7, stride=1, pad=3, pad_type='reflect', scope='G_logit')
+            x = tanh(x)
+
+            return x, cam_logit, heatmap
+
+    def MLP(self, x, use_bias=True, reuse=False, scope='MLP'):
+        channel = self.ch * self.n_res
+
+        if self.light :
+            x = global_avg_pooling(x)
+
+        with tf.variable_scope(scope, reuse=reuse):
+            for i in range(2) :
+                x = fully_connected(x, channel, use_bias, scope='linear_' + str(i))
+                x = relu(x)
+
+
+            gamma = fully_connected(x, channel, use_bias, scope='gamma')
+            beta = fully_connected(x, channel, use_bias, scope='beta')
+
+            gamma = tf.reshape(gamma, shape=[self.batch_size, 1, 1, channel])
+            beta = tf.reshape(beta, shape=[self.batch_size, 1, 1, channel])
+
+            return gamma, beta
+
+    ##################################################################################
+    # Discriminator
+    ##################################################################################
+
+    def discriminator(self, x_init, reuse=False, scope="discriminator"):
+        D_logit = []
+        D_CAM_logit = []
+        with tf.variable_scope(scope, reuse=reuse) :
+            local_x, local_cam, local_heatmap = self.discriminator_local(x_init, reuse=reuse, scope='local')
+            global_x, global_cam, global_heatmap = self.discriminator_global(x_init, reuse=reuse, scope='global')
+
+            D_logit.extend([local_x, global_x])
+            D_CAM_logit.extend([local_cam, global_cam])
+
+            return D_logit, D_CAM_logit, local_heatmap, global_heatmap
+
+    def discriminator_global(self, x_init, reuse=False, scope='discriminator_global'):
+        with tf.variable_scope(scope, reuse=reuse):
+            channel = self.ch
+            x = conv(x_init, channel, kernel=4, stride=2, pad=1, pad_type='reflect', sn=self.sn, scope='conv_0')
+            x = lrelu(x, 0.2)
+
+            for i in range(1, self.n_dis - 1):
+                x = conv(x, channel * 2, kernel=4, stride=2, pad=1, pad_type='reflect', sn=self.sn, scope='conv_' + str(i))
+                x = lrelu(x, 0.2)
+
+                channel = channel * 2
+
+            x = conv(x, channel * 2, kernel=4, stride=1, pad=1, pad_type='reflect', sn=self.sn, scope='conv_last')
+            x = lrelu(x, 0.2)
+
+            channel = channel * 2
+
+            cam_x = global_avg_pooling(x)
+            cam_gap_logit, cam_x_weight = fully_connected_with_w(cam_x, sn=self.sn, scope='CAM_logit')
+            x_gap = tf.multiply(x, cam_x_weight)
+
+            cam_x = global_max_pooling(x)
+            cam_gmp_logit, cam_x_weight = fully_connected_with_w(cam_x, sn=self.sn, reuse=True, scope='CAM_logit')
+            x_gmp = tf.multiply(x, cam_x_weight)
+
+            cam_logit = tf.concat([cam_gap_logit, cam_gmp_logit], axis=-1)
+            x = tf.concat([x_gap, x_gmp], axis=-1)
+
+            x = conv(x, channel, kernel=1, stride=1, scope='conv_1x1')
+            x = lrelu(x, 0.2)
+
+            heatmap = tf.squeeze(tf.reduce_sum(x, axis=-1))
+
+
+            x = conv(x, channels=1, kernel=4, stride=1, pad=1, pad_type='reflect', sn=self.sn, scope='D_logit')
+
+            return x, cam_logit, heatmap
+
+    def discriminator_local(self, x_init, reuse=False, scope='discriminator_local'):
+        with tf.variable_scope(scope, reuse=reuse) :
+            channel = self.ch
+            x = conv(x_init, channel, kernel=4, stride=2, pad=1, pad_type='reflect', sn=self.sn, scope='conv_0')
+            x = lrelu(x, 0.2)
+
+            for i in range(1, self.n_dis - 2 - 1):
+                x = conv(x, channel * 2, kernel=4, stride=2, pad=1, pad_type='reflect', sn=self.sn, scope='conv_' + str(i))
+                x = lrelu(x, 0.2)
+
+                channel = channel * 2
+
+            x = conv(x, channel * 2, kernel=4, stride=1, pad=1, pad_type='reflect', sn=self.sn, scope='conv_last')
+            x = lrelu(x, 0.2)
+
+            channel = channel * 2
+
+            cam_x = global_avg_pooling(x)
+            cam_gap_logit, cam_x_weight = fully_connected_with_w(cam_x, sn=self.sn, scope='CAM_logit')
+            x_gap = tf.multiply(x, cam_x_weight)
+
+            cam_x = global_max_pooling(x)
+            cam_gmp_logit, cam_x_weight = fully_connected_with_w(cam_x, sn=self.sn, reuse=True, scope='CAM_logit')
+            x_gmp = tf.multiply(x, cam_x_weight)
+
+            cam_logit = tf.concat([cam_gap_logit, cam_gmp_logit], axis=-1)
+            x = tf.concat([x_gap, x_gmp], axis=-1)
+
+            x = conv(x, channel, kernel=1, stride=1, scope='conv_1x1')
+            x = lrelu(x, 0.2)
+
+            heatmap = tf.squeeze(tf.reduce_sum(x, axis=-1))
+
+            x = conv(x, channels=1, kernel=4, stride=1, pad=1, pad_type='reflect', sn=self.sn, scope='D_logit')
+
+            return x, cam_logit, heatmap
+
+    def generate_a2b(self, x_A, reuse=False):
+        out, cam, _ = self.generator(x_A, reuse=reuse, scope="generator_B")
+
+        return out, cam
+
+    def generate_b2a(self, x_B, reuse=False):
+        out, cam, _ = self.generator(x_B, reuse=reuse, scope="generator_A")
+
+        return out, cam
+    def build_model(self):
+        self.test_domain_A = tf.placeholder(tf.float32, [1, self.img_size, self.img_size, self.img_ch], name='test_domain_A')
+        self.test_domain_B = tf.placeholder(tf.float32, [1, self.img_size, self.img_size, self.img_ch], name='test_domain_B')
+
+        self.test_fake_B, _ = self.generate_a2b(self.test_domain_A)
+        self.test_fake_A, _ = self.generate_b2a(self.test_domain_B)
+
+    @property
+    def model_dir(self):
+        n_res = str(self.n_res) + 'resblock'
+        n_dis = str(self.n_dis) + 'dis'
+
+        if self.smoothing:
+            smoothing = '_smoothing'
+        else:
+            smoothing = ''
+
+        if self.sn:
+            sn = '_sn'
+        else:
+            sn = ''
+
+        return "{}_{}_{}_{}_{}_{}_{}_{}_{}_{}{}{}".format(self.model_name, self.dataset_name,
+                                                          self.gan_type, n_res, n_dis,
+                                                          self.n_critic,
+                                                          self.adv_weight, self.cycle_weight, self.identity_weight,
+                                                          self.cam_weight, sn, smoothing)
+
+    def load(self, checkpoint_dir):
+        print(" [*] Reading checkpoints...")
+        checkpoint_dir = os.path.join(checkpoint_dir, self.model_dir)
+
+        ckpt = tf.train.get_checkpoint_state(checkpoint_dir)
+        if ckpt and ckpt.model_checkpoint_path:
+            ckpt_name = os.path.basename(ckpt.model_checkpoint_path)
+            self.saver.restore(self.sess, os.path.join(checkpoint_dir, ckpt_name))
+            counter = int(ckpt_name.split('-')[-1])
+            print(" [*] Success to read {}".format(ckpt_name))
+            return True, counter
+        else:
+            print(" [*] Failed to find a checkpoint")
+            return False, 0
+
+    def loadModel(self):
+        tf.global_variables_initializer().run(session=self.sess)
+
+        self.saver = tf.train.Saver()
+        could_load, checkpoint_counter = self.load(self.checkpoint_dir)
+        self.result_dir = os.path.join(self.result_dir, self.model_dir)
+        check_folder(self.result_dir)
+
+        if could_load:
+            print(" [*] Load SUCCESS")
+        else:
+            print(" [!] Load failed...")
+
+    def test(self, sample_file):
+         # A -> B
+        print('Processing A image: ' + sample_file)
+        sample_image = np.asarray(load_test_data(sample_file, size=self.img_size))
+        image_path = os.path.join(self.result_dir,'{0}'.format(os.path.basename(sample_file)))
+
+        fake_img = self.sess.run(self.test_fake_B, feed_dict = {self.test_domain_A : sample_image})
+        save_images(fake_img, [1, 1], image_path)
+
+        return image_path
+
+
+gan = None
+def main_test(img_path):
+    # open session
+    sess = tf.Session(config=tf.ConfigProto(allow_soft_placement=True))
+    global gan
+    if gan is None:
+        gan = UgatitTest(sess)
+        # build graph
+        gan.build_model()
+        # show network architecture
+        show_all_variables()
+
+        gan.loadModel()
+
+    result = gan.test(img_path)
+    print(" [*] Test finished!")
+    print(result)
+    return os.path.abspath(result)
+
+if __name__ == '__main__':
+    main_test('/home/hylee/cartoon/myp2c/imgs/src/im4.jpg')