Upload 28 files

.gitattributes

@@ -33,3 +33,11 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+assets/augan_result.png filter=lfs diff=lfs merge=lfs -text
+assets/augan_uncer.png filter=lfs diff=lfs merge=lfs -text
+datasets/swim/testA/GP010594_frame_000017_rgb_anon.png filter=lfs diff=lfs merge=lfs -text
+datasets/swim/testA/GP010594_frame_000021_rgb_anon.png filter=lfs diff=lfs merge=lfs -text
+datasets/swim/testA/GP010594_frame_000087_rgb_anon.png filter=lfs diff=lfs merge=lfs -text
+datasets/swim/testB/GOPR0351_frame_000159_rgb_ref_anon.png filter=lfs diff=lfs merge=lfs -text
+datasets/swim/testB/GOPR0351_frame_000161_rgb_ref_anon.png filter=lfs diff=lfs merge=lfs -text
+datasets/swim/testB/GOPR0355_frame_000138_rgb_ref_anon.png filter=lfs diff=lfs merge=lfs -text

AUGAN.py

@@ -0,0 +1,738 @@
+from collections import namedtuple
+from models import generator_resnet, discriminator
+from utils import *
+from loss_utils import *
+from ops import *
+import time
+import matplotlib.pyplot as plt
+from glob import glob
+
+
+class AUGAN(object):
+    def __init__(self, sess, args):
+        self.sess = sess
+        self.batch_size = args.batch_size
+        self.image_size = args.fine_size
+        self.input_c_dim = args.input_nc
+        self.output_c_dim = args.output_nc
+        self.L1_lambda = args.L1_lambda
+        self.conf_lambda = args.conf_lambda
+        self.dataset_dir = args.dataset_dir
+        self.n_d = args.n_d
+        self.n_scale = args.n_scale
+        self.ndf = args.ndf
+        self.load_size = args.load_size
+        self.fine_size = args.fine_size
+        self.generator = generator_resnet
+        self.discriminator = discriminator
+        if args.use_lsgan:
+            self.criterionGAN = mae_criterion
+            self.criterionGAN_list = mae_criterion_list
+        else:
+            self.criterionGAN = sce_criterion
+            self.criterionGAN_list = sce_criterion_list
+
+        self.use_uncertainty = args.use_uncertainty
+
+        OPTIONS = namedtuple(
+            "OPTIONS",
+            "batch_size image_size \
+                              gf_dim df_dim output_c_dim is_training",
+        )
+        self.options = OPTIONS._make(
+            (
+                args.batch_size,
+                args.fine_size,
+                args.ngf,
+                args.ndf // args.n_d,
+                args.output_nc,
+                args.phase == "train",
+            )
+        )
+        self.save_conf = args.save_conf
+        self._build_model()
+        self.saver = tf.compat.v1.train.Saver()
+        self.pool = ImagePool(args.max_size)
+
+    def _build_model(self):
+        self.real_data = tf.compat.v1.placeholder(
+            tf.float32,
+            [
+                self.batch_size,
+                self.image_size,
+                self.image_size * 2,
+                self.input_c_dim + self.output_c_dim,
+            ],
+            name="real_A_and_B_images",
+        )
+
+        self.real_A = self.real_data[:, :, :, : self.input_c_dim]
+        self.real_B = self.real_data[
+            :, :, :, self.input_c_dim : self.input_c_dim + self.output_c_dim
+        ]
+
+        A_label = np.zeros([1, 1, 1, 2], dtype=np.float32)
+        B_label = np.zeros([1, 1, 1, 2], dtype=np.float32)
+        A_label[:, :, :, 0] = 1.0
+        B_label[:, :, :, 1] = 1.0
+        self.A_label = tf.convert_to_tensor(A_label)
+        self.B_label = tf.convert_to_tensor(B_label)
+
+        (
+            self.fake_B,
+            self.rec_realA,
+            self.realA_percep,
+            self.transA_percep,
+            self.pred_confA,
+        ) = self.generator(
+            self.real_A, self.options, transfer=True, reuse=False, name="generatorA2B"
+        )
+        self.fake_A_, self.rec_fakeB, self.fakeB_percep, _, _ = self.generator(
+            self.fake_B, self.options, transfer=False, reuse=False, name="generatorB2A"
+        )
+        self.fake_A, self.rec_realB, self.realB_percep, _, _ = self.generator(
+            self.real_B, self.options, transfer=False, reuse=True, name="generatorB2A"
+        )
+        self.fake_B_, self.rec_fakeA, self.fakeA_percep, self.trans_fakeA_percep, _ = (
+            self.generator(
+                self.fake_A,
+                self.options,
+                transfer=True,
+                reuse=True,
+                name="generatorA2B",
+            )
+        )
+
+        self.g_adv_total = 0.0
+        self.g_adv = 0.0
+        self.g_adv_rec = 0.0
+        self.g_adv_recfake = 0.0
+
+        self.percep_loss = tf.reduce_mean(
+            tf.abs(
+                tf.reduce_mean(self.transA_percep, axis=3)
+                - tf.reduce_mean(self.fakeB_percep, axis=3)
+            )
+        ) + tf.reduce_mean(
+            tf.abs(
+                tf.reduce_mean(self.realB_percep, axis=3)
+                - tf.reduce_mean(self.fakeA_percep, axis=3)
+            )
+        )
+
+        for i in range(self.n_d):
+            self.DB_fake = self.discriminator(
+                self.fake_B, self.options, reuse=False, name=str(i) + "_discriminatorB"
+            )
+            self.DA_fake = self.discriminator(
+                self.fake_A, self.options, reuse=False, name=str(i) + "_discriminatorA"
+            )
+
+            self.g_adv_total += self.criterionGAN_list(
+                self.DA_fake, get_ones_like(self.DA_fake)
+            ) + self.criterionGAN_list(self.DB_fake, get_ones_like(self.DB_fake))
+
+            self.g_adv += self.criterionGAN_list(
+                self.DA_fake, get_ones_like(self.DA_fake)
+            ) + self.criterionGAN_list(self.DB_fake, get_ones_like(self.DB_fake))
+
+        self.g_loss_a2b = (
+            self.criterionGAN_list(self.DB_fake, get_ones_like(self.DB_fake))
+            + self.L1_lambda * abs_criterion(self.real_A, self.fake_A_)
+            + self.L1_lambda * abs_criterion(self.real_B, self.fake_B_)
+        )
+        self.g_loss_b2a = (
+            self.criterionGAN_list(self.DA_fake, get_ones_like(self.DA_fake))
+            + self.L1_lambda * abs_criterion(self.real_A, self.fake_A_)
+            + self.L1_lambda * abs_criterion(self.real_B, self.fake_B_)
+        )
+
+        self.g_A_recon_loss = self.L1_lambda * abs_criterion(
+            self.rec_realA, self.real_A
+        )
+        self.g_B_recon_loss = self.L1_lambda * abs_criterion(
+            self.rec_realB, self.real_B
+        )
+        if self.use_uncertainty:
+            self.g_A_cycle_loss = self.conf_lambda * conf_criterion(
+                self.real_A, self.fake_A_, self.pred_confA
+            )
+        else:
+            self.g_A_cycle_loss = self.L1_lambda * abs_criterion(
+                self.real_A, self.fake_A_
+            )
+        self.g_B_cylce_loss = self.L1_lambda * abs_criterion(self.real_B, self.fake_B_)
+
+        self.g_loss = (
+            self.g_adv_total
+            + self.g_A_recon_loss
+            + self.g_B_recon_loss
+            + self.g_A_cycle_loss
+            + self.g_B_cylce_loss
+            + self.percep_loss
+        )
+
+        self.g_rec_real = abs_criterion(self.rec_realA, self.real_A) + abs_criterion(
+            self.rec_realB, self.real_B
+        )
+        self.g_rec_cycle = abs_criterion(self.real_A, self.fake_A_) + abs_criterion(
+            self.real_B, self.fake_B_
+        )
+
+        self.fake_A_sample = tf.compat.v1.placeholder(
+            tf.float32,
+            [self.batch_size, self.image_size, self.image_size * 2, self.output_c_dim],
+            name="fake_A_sample",
+        )
+        self.fake_B_sample = tf.compat.v1.placeholder(
+            tf.float32,
+            [self.batch_size, self.image_size, self.image_size * 2, self.output_c_dim],
+            name="fake_B_sample",
+        )
+        self.rec_A_sample = tf.compat.v1.placeholder(
+            tf.float32,
+            [self.batch_size, self.image_size, self.image_size * 2, self.output_c_dim],
+            name="rec_A_sample",
+        )
+        self.rec_B_sample = tf.compat.v1.placeholder(
+            tf.float32,
+            [self.batch_size, self.image_size, self.image_size * 2, self.output_c_dim],
+            name="rec_B_sample",
+        )
+        self.rec_fakeA_sample = tf.compat.v1.placeholder(
+            tf.float32,
+            [self.batch_size, self.image_size, self.image_size * 2, self.output_c_dim],
+            name="rec_fakeA_sample",
+        )
+        self.rec_fakeB_sample = tf.compat.v1.placeholder(
+            tf.float32,
+            [self.batch_size, self.image_size, self.image_size * 2, self.output_c_dim],
+            name="rec_fakeB_sample",
+        )
+
+        self.d_loss_item = []
+        self.d_loss_item_rec = []
+        self.d_loss_item_recfake = []
+
+        for i in range(self.n_d):
+            self.DB_real = self.discriminator(
+                self.real_B, self.options, reuse=True, name=str(i) + "_discriminatorB"
+            )
+            self.DA_real = self.discriminator(
+                self.real_A, self.options, reuse=True, name=str(i) + "_discriminatorA"
+            )
+            self.DB_fake_sample = self.discriminator(
+                self.fake_B_sample,
+                self.options,
+                reuse=True,
+                name=str(i) + "_discriminatorB",
+            )
+            self.DA_fake_sample = self.discriminator(
+                self.fake_A_sample,
+                self.options,
+                reuse=True,
+                name=str(i) + "_discriminatorA",
+            )
+            self.db_loss_real = self.criterionGAN_list(
+                self.DB_real, get_ones_like(self.DB_real)
+            )
+            self.db_loss_fake = self.criterionGAN_list(
+                self.DB_fake_sample, get_zeros_like(self.DB_fake_sample)
+            )
+            self.db_loss = self.db_loss_real * 0.5 + self.db_loss_fake * 0.5
+            self.da_loss_real = self.criterionGAN_list(
+                self.DA_real, get_ones_like(self.DA_real)
+            )
+            self.da_loss_fake = self.criterionGAN_list(
+                self.DA_fake_sample, get_zeros_like(self.DA_fake_sample)
+            )
+            self.da_loss = self.da_loss_real * 0.5 + self.da_loss_fake * 0.5
+            self.d_loss = self.da_loss + self.db_loss
+            self.d_loss_item.append(self.d_loss)
+
+        self.g_loss_a2b_sum = tf.compat.v1.summary.scalar("g_loss_a2b", self.g_loss_a2b)
+        self.g_loss_b2a_sum = tf.compat.v1.summary.scalar("g_loss_b2a", self.g_loss_b2a)
+        self.g_loss_sum = tf.compat.v1.summary.scalar("g_loss", self.g_loss)
+        self.g_sum = tf.compat.v1.summary.merge(
+            [self.g_loss_a2b_sum, self.g_loss_b2a_sum, self.g_loss_sum]
+        )
+        self.db_loss_sum = tf.compat.v1.summary.scalar("db_loss", self.db_loss)
+        self.da_loss_sum = tf.compat.v1.summary.scalar("da_loss", self.da_loss)
+        self.d_loss_sum = tf.compat.v1.summary.scalar("d_loss", self.d_loss)
+        self.db_loss_real_sum = tf.compat.v1.summary.scalar(
+            "db_loss_real", self.db_loss_real
+        )
+        self.db_loss_fake_sum = tf.compat.v1.summary.scalar(
+            "db_loss_fake", self.db_loss_fake
+        )
+        self.da_loss_real_sum = tf.compat.v1.summary.scalar(
+            "da_loss_real", self.da_loss_real
+        )
+        self.da_loss_fake_sum = tf.compat.v1.summary.scalar(
+            "da_loss_fake", self.da_loss_fake
+        )
+        self.d_sum = tf.compat.v1.summary.merge(
+            [
+                self.da_loss_sum,
+                self.da_loss_real_sum,
+                self.da_loss_fake_sum,
+                self.db_loss_sum,
+                self.db_loss_real_sum,
+                self.db_loss_fake_sum,
+                self.d_loss_sum,
+            ]
+        )
+
+        self.test_A = tf.compat.v1.placeholder(
+            tf.float32,
+            [self.batch_size, self.image_size, self.image_size * 2, self.input_c_dim],
+            name="test_A",
+        )
+        self.test_B = tf.compat.v1.placeholder(
+            tf.float32,
+            [self.batch_size, self.image_size, self.image_size * 2, self.output_c_dim],
+            name="test_B",
+        )
+
+        (
+            self.testB,
+            self.rec_testA,
+            self.testA_percep,
+            self.trans_testA_percep,
+            self.test_pred_confA,
+        ) = self.generator(
+            self.test_A, self.options, transfer=True, reuse=True, name="generatorA2B"
+        )
+        self.rec_cycle_A, self.refine_testB, self.testB_percep, _, _ = self.generator(
+            self.testB, self.options, transfer=False, reuse=True, name="generatorB2A"
+        )
+
+        self.testA, self.rec_testB, _, _, _ = self.generator(
+            self.test_B, self.options, transfer=False, reuse=True, name="generatorB2A"
+        )
+        self.rec_cycle_B, self.refine_testA, _, _, _ = self.generator(
+            self.testA, self.options, True, True, name="generatorA2B"
+        )
+
+        t_vars = tf.compat.v1.trainable_variables()
+
+        self.g_vars = [var for var in t_vars if "generator" in var.name]
+        self.p_vars = [var for var in t_vars if "percep" in var.name]
+        self.d_vars_item = []
+        for i in range(self.n_d):
+            self.d_vars = [
+                var for var in t_vars if str(i) + "_discriminator" in var.name
+            ]
+            self.d_vars_item.append(self.d_vars)
+
+    def train(self, args):
+
+        self.lr = tf.compat.v1.placeholder(tf.float32, None, name="learning_rate")
+
+        ### generator
+        self.g_optim = tf.optimizers.Adam(
+            learning_rate=self.lr, beta_1=args.beta1
+        ).minimize(self.g_loss, var_list=self.g_vars, tape=None)
+
+        ### translation
+        self.d_optim_item = []
+        for i in range(self.n_d):
+            self.d_optim = tf.optimizers.Adam(
+                learning_rate=self.lr, beta_1=args.beta1
+            ).minimize(self.g_loss, var_list=self.g_vars, tape=None)
+            self.d_optim_item.append(self.d_optim)
+
+        init_op = tf.compat.v1.global_variables_initializer()
+        self.sess.run(init_op)
+        self.writer = tf.summary.FileWriter(
+            os.path.join(args.checkpoint_dir, "logs"), self.sess.graph
+        )
+
+        counter = 1
+        start_time = time.time()
+
+        if args.continue_train:
+            if self.load(args.checkpoint_dir):
+                print(" [*] Load SUCCESS")
+            else:
+                print(" [!] Load failed...")
+
+        print("Training.........................")
+        for epoch in range(args.epoch):
+            dataA = glob("./datasets/{}/*.*".format(self.dataset_dir + "/trainA"))
+            dataB = glob("./datasets/{}/*.*".format(self.dataset_dir + "/trainB"))
+            if (len(dataA) == 0) or (len(dataB) == 0):
+                raise Exception("No files found in the dataset")
+            else:
+                print(
+                    "Data found in the dataset. length of A: ",
+                    len(dataA),
+                    " B: ",
+                    len(dataB),
+                )
+            np.random.shuffle(dataA)
+            np.random.shuffle(dataB)
+            batch_idxs = (
+                min(min(len(dataA), len(dataB)), args.train_size) // self.batch_size
+            )
+            lr = (
+                args.lr
+                if epoch < args.epoch_step
+                else args.lr * (args.epoch - epoch) / (args.epoch - args.epoch_step)
+            )
+
+            for idx in range(0, batch_idxs):
+                print("Epoch: [%2d] [%4d/%4d] " % (epoch, idx, batch_idxs))
+                batch_files = list(
+                    zip(
+                        dataA[idx * self.batch_size : (idx + 1) * self.batch_size],
+                        dataB[idx * self.batch_size : (idx + 1) * self.batch_size],
+                    )
+                )
+                batch_images = [
+                    load_train_data(batch_file, args.load_size, args.fine_size)
+                    for batch_file in batch_files
+                ]
+                batch_images = np.array(batch_images).astype(np.float32)
+                # Update G network and record fake outputs
+                print("Training G network----------------------")
+                (
+                    fake_A,
+                    fake_B,
+                    rec_A,
+                    rec_B,
+                    rec_fake_A,
+                    rec_fake_B,
+                    _,
+                    g_loss,
+                    gan_loss,
+                    percep,
+                    g_adv,
+                    g_A_recon_loss,
+                    g_B_recon_loss,
+                    g_A_cycle_loss,
+                    g_B_cycle_loss,
+                    summary_str,
+                ) = self.sess.run(
+                    [
+                        self.fake_A,
+                        self.fake_B,
+                        self.rec_realA,
+                        self.rec_realB,
+                        self.rec_fakeA,
+                        self.rec_fakeB,
+                        self.g_optim,
+                        self.g_loss,
+                        self.g_adv_total,
+                        self.percep_loss,
+                        self.g_adv,
+                        self.g_A_recon_loss,
+                        self.g_B_recon_loss,
+                        self.g_A_cycle_loss,
+                        self.g_B_cylce_loss,
+                        self.g_sum,
+                    ],
+                    feed_dict={self.real_data: batch_images, self.lr: lr},
+                )
+                self.writer.add_summary(summary_str, counter)
+                [fake_A, fake_B] = self.pool([fake_A, fake_B])
+
+                # Update D network
+                print("Training D network----------------------")
+                loss_print = []
+                for i in range(self.n_d):
+                    _, d_loss, d_sum = self.sess.run(
+                        [self.d_optim_item[i], self.d_loss_item[i], self.d_sum],
+                        feed_dict={
+                            self.real_data: batch_images,
+                            self.fake_A_sample: fake_A,
+                            self.fake_B_sample: fake_B,
+                            self.lr: lr,
+                        },
+                    )
+
+                    loss_print.append(d_loss)
+
+                counter += 1
+                print(
+                    (
+                        "Epoch: [%2d] [%4d/%4d] time: %4.4f g_loss: %4.4f gan:%4.4f adv:%4.4f g_percep:%4.4f "
+                        % (
+                            epoch,
+                            idx,
+                            batch_idxs,
+                            time.time() - start_time,
+                            g_loss,
+                            gan_loss,
+                            g_adv,
+                            percep,
+                        )
+                    )
+                )
+
+                if np.mod(counter, args.print_freq) == 1:
+                    self.sample_model(args.sample_dir, epoch, idx)
+
+                if np.mod(counter, args.save_freq) == 2:
+                    self.save(args.checkpoint_dir, counter)
+
+    def save(self, checkpoint_dir, step):
+        model_name = "cyclegan.model"
+        model_dir = "%s_%s" % (self.dataset_dir, self.image_size)
+        checkpoint_dir = os.path.join(checkpoint_dir, model_dir)
+
+        if not os.path.exists(checkpoint_dir):
+            os.makedirs(checkpoint_dir)
+
+        self.saver.save(
+            self.sess, os.path.join(checkpoint_dir, model_name), global_step=step
+        )
+
+    def load(self, checkpoint_dir):
+        print(" [*] Reading checkpoint...")
+
+        model_dir = "%s_%s" % (self.dataset_dir, self.image_size)
+        checkpoint_dir = os.path.join(checkpoint_dir, 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))
+            return True
+        else:
+            return False
+
+    def sample_model(self, sample_dir, epoch, idx):
+        dataA = glob("./datasets/{}/*.*".format(self.dataset_dir + "/testA"))
+        dataB = glob("./datasets/{}/*.*".format(self.dataset_dir + "/testB"))
+        if (len(dataA) == 0) or (len(dataB) == 0):
+            raise Exception("No files found in the test directory")
+        np.random.shuffle(dataA)
+        np.random.shuffle(dataB)
+        batch_files = list(zip(dataA[: self.batch_size], dataB[: self.batch_size]))
+        sample_images = [
+            load_train_data(batch_file, self.load_size, self.fine_size, is_testing=True)
+            for batch_file in batch_files
+        ]
+        sample_images = np.array(sample_images).astype(np.float32)
+
+        fake_A, fake_B = self.sess.run(
+            [self.fake_A, self.fake_B], feed_dict={self.real_data: sample_images}
+        )
+        real_A = sample_images[:, :, :, :3]
+        real_B = sample_images[:, :, :, 3:]
+
+        merge_A = np.concatenate([real_B, fake_A], axis=2)
+        merge_B = np.concatenate([real_A, fake_B], axis=2)
+        check_folder("./{}/{:02d}".format(sample_dir, epoch))
+        save_images(
+            merge_A,
+            [self.batch_size, 1],
+            "./{}/{:02d}/A_{:04d}.jpg".format(sample_dir, epoch, idx),
+        )
+        save_images(
+            merge_B,
+            [self.batch_size, 1],
+            "./{}/{:02d}/B_{:04d}.jpg".format(sample_dir, epoch, idx),
+        )
+
+    def test(self, args):
+        total_time = 0
+
+        init_op = tf.compat.v1.global_variables_initializer()
+        self.sess.run(init_op)
+        if args.which_direction == "AtoB":
+            sample_files = glob("./datasets/{}/*.*".format(self.dataset_dir + "/testA"))
+        elif args.which_direction == "BtoA":
+            sample_files = glob("./datasets/{}/*.*".format(self.dataset_dir + "/testB"))
+        else:
+            raise Exception("--which_direction must be AtoB or BtoA")
+
+        if len(sample_files) == 0:
+            raise Exception("No files found in the test directory")
+
+        # print(sample_files)
+
+        if self.load(args.checkpoint_dir):
+            print(" [*] Load SUCCESS")
+        else:
+            print(" [!] Load failed...")
+        out_var, refine_var, in_var, rec_var, cycle_var, percep_var, conf_var = (
+            (
+                self.testB,
+                self.refine_testB,
+                self.test_A,
+                self.rec_testA,
+                self.rec_cycle_A,
+                self.testA_percep,
+                self.test_pred_confA,
+            )
+            if args.which_direction == "AtoB"
+            else (
+                self.testA,
+                self.refine_testA,
+                self.test_B,
+                self.rec_testB,
+                self.rec_cycle_B,
+                self.testB_percep,
+                self.test_pred_confA,
+            )
+        )
+        for sample_file in sample_files:
+            # print('Processing image: ' + sample_file)
+            sample_image = [load_test_data(sample_file, args.fine_size)]
+            start_time = time.time()
+            sample_image = np.array(sample_image).astype(np.float32)
+            image_path = os.path.join(
+                args.test_dir,
+                "{0}_{1}".format(args.which_direction, os.path.basename(sample_file)),
+            )
+            ori_path = os.path.join(
+                args.test_dir,
+                "{0}_{1}".format("ori", os.path.basename(sample_file)),
+            )
+            conf_path = os.path.join(
+                args.conf_dir,
+                "{0}_{1}".format(args.which_direction, os.path.basename(sample_file)),
+            )
+
+            (fake_img,) = self.sess.run([out_var], feed_dict={in_var: sample_image})
+            end_time = time.time()
+            # merge = np.concatenate([sample_image, fake_img], axis=2)
+            save_images(fake_img[0], [1], image_path)
+            save_images(sample_image[0], [1], ori_path)
+            # save_images(merge, [1, 1], image_path)
+            total_time = total_time + (end_time - start_time)
+
+            if args.save_conf:
+
+                if args.which_direction == "AtoB":
+                    pass
+                else:
+                    raise Exception(
+                        "--conf map only can be estimated in AtoB direction"
+                    )
+
+                conf_img = self.sess.run(conf_var, feed_dict={in_var: sample_image})
+                conf_img_sq = np.squeeze(conf_img)
+                plt.imshow(
+                    conf_img_sq, cmap="plasma", interpolation="nearest", alpha=1.0
+                )
+                plt.savefig(conf_path)
+        print(
+            f"Average time taken to convert images: {total_time/len(sample_files)} seconds"
+        )
+
+    def convert(self, args, datadir="./inf_data"):
+        total_time = 0
+
+        init_op = tf.compat.v1.global_variables_initializer()
+        self.sess.run(init_op)
+
+        if self.load(args.checkpoint_dir):
+            print(" [*] Load SUCCESS")
+        else:
+            raise Exception("-- Cannot Load Model. Train or Add model first")
+
+        if args.which_direction == "AtoB":
+            sample_files = glob(datadir)
+        elif args.which_direction == "BtoA":
+            sample_files = glob(datadir)
+        else:
+            raise Exception("--which_direction must be AtoB or BtoA")
+
+        print(sample_files)
+
+        out_var, refine_var, in_var, rec_var, cycle_var, percep_var, conf_var = (
+            (
+                self.testB,
+                self.refine_testB,
+                self.test_A,
+                self.rec_testA,
+                self.rec_cycle_A,
+                self.testA_percep,
+                self.test_pred_confA,
+            )
+            if args.which_direction == "AtoB"
+            else (
+                self.testA,
+                self.refine_testA,
+                self.test_B,
+                self.rec_testB,
+                self.rec_cycle_B,
+                self.testB_percep,
+                self.test_pred_confA,
+            )
+        )
+        for sample_file in sample_files:
+            print("Processing image: " + sample_file)
+            sample_image = [load_test_data(sample_file, args.fine_size)]
+            start_time = time.time()
+            sample_image = np.array(sample_image).astype(np.float32)
+            image_path = os.path.join(
+                args.test_dir,
+                "{0}_{1}".format(args.which_direction, os.path.basename(sample_file)),
+            )
+            conf_path = os.path.join(
+                args.conf_dir,
+                "{0}_{1}".format(args.which_direction, os.path.basename(sample_file)),
+            )
+
+            (fake_img,) = self.sess.run([out_var], feed_dict={in_var: sample_image})
+            end_time = time.time()
+            merge = np.concatenate([sample_image, fake_img], axis=2)
+            save_images(merge, [1, 1], image_path)
+            total_time = total_time + (end_time - start_time)
+            print(f"Time taken to convert image: {end_time - start_time} seconds")
+
+            if args.save_conf:
+
+                if args.which_direction == "AtoB":
+                    pass
+                else:
+                    raise Exception(
+                        "--conf map only can be estimated in AtoB direction"
+                    )
+
+                conf_img = self.sess.run(conf_var, feed_dict={in_var: sample_image})
+                conf_img_sq = np.squeeze(conf_img)
+                plt.imshow(
+                    conf_img_sq, cmap="plasma", interpolation="nearest", alpha=1.0
+                )
+                plt.savefig(conf_path)
+        print(
+            f"Average time taken to convert images: {total_time/len(sample_files)} seconds"
+        )
+
+    def convert_image(self, args, input_image_path, output_dir):
+        init_op = tf.compat.v1.global_variables_initializer()
+        if self.load(args.checkpoint_dir):
+            print(" [*] Load SUCCESS")
+        with tf.Session() as sess:
+            sess.run(init_op)
+            # Load the input image
+            input_image = [load_test_data(input_image_path, self.fine_size)]
+            input_image = np.array(input_image).astype(np.float32)
+
+            # Get the generator output
+            if args.which_direction == "AtoB":
+                out_var = self.testB
+                in_var = self.test_A
+            else:
+                out_var = self.testA
+                in_var = self.test_B
+
+            # Run the model to obtain the converted image
+            start_time = time.time()
+            converted_image = sess.run(out_var, feed_dict={in_var: input_image})
+            end_time = time.time()
+
+            # Save the converted image
+            output_image_path = os.path.join(
+                output_dir, os.path.basename(input_image_path)
+            )
+            merge = np.concatenate([input_image, converted_image], axis=2)
+            save_images(merge, [1, 1], output_image_path)
+
+            # Print the time taken
+            print(f"Time taken to convert image: {end_time - start_time} seconds")

README.md

@@ -0,0 +1,125 @@
+# Adverse Weather Image Translation with Asymmetric and Uncertainty-aware GAN (AU-GAN)
+Official Tensorflow implementation of [Adverse Weather Image Translation with Asymmetric and Uncertainty-aware GAN](https://www.bmvc2021-virtualconference.com/assets/papers/1443.pdf) (AU-GAN)\
+Jeong-gi Kwak, Youngsaeng Jin, Yuanming Li, Dongsik Yoon, Donghyeon Kim and Hanseok Ko </br>
+*British Machine Vision Conference (BMVC), 2021*
+</br>
+
+## Intro 
+
+### Night &rarr; Day ([BDD100K](https://bdd-data.berkeley.edu/))
+<img src="./assets/augan_bdd.png" width="800">
+
+### Rainy night &rarr; Day ([Alderdey](https://wiki.qut.edu.au/pages/viewpage.action?pageId=181178395))
+<img src="./assets/augan_alderley.png" width="800">
+</br>
+
+
+## Architecture
+<img src="./assets/augan_model.png" width="800">
+Our generator has asymmetric structure for editing day&rarr;night and night&rarr;day.
+Please refer our paper for details
+
+## **Envs**
+
+```bash
+
+git clone https://github.com/jgkwak95/AU-GAN.git
+cd AU-GAN
+
+# Create virtual environment
+conda create -y --name augan python=3.6.7
+conda activate augan
+
+conda install tensorflow-gpu==1.14.0   # Tensorflow 1.14
+pip install --no-cache-dir -r requirements.txt
+
+```
+
+## **Preparing datasets**
+
+**Night &rarr; Day** </br>
+[Berkeley DeepDrive dataset](https://bdd-data.berkeley.edu/) contains 100,000 high resolution images of the urban roads for autonomous driving.</br></br>
+**Rainy night &rarr; Day** </br>
+[Alderley dataset](https://wiki.qut.edu.au/pages/viewpage.action?pageId=181178395) consists of images of two domains,
+rainy night and daytime. It was collected while driving the same route in each weather environment.</br>
+</br>
+Please download datasets and then construct them following [ForkGAN](https://github.com/zhengziqiang/ForkGAN)
+
+## Pretrained Model 
+
+Download the pretrained model for BDD100K(256x512) [here](https://drive.google.com/file/d/1rvIF3yE9MwPWj0kD4IEstETyMQXYAHzr/view?usp=sharing) and unzip it to ./check/bdd_exp/bdd100k_256/
+
+## Training
+
+```bash
+
+# Alderley (256x512)
+python main_uncer.py --dataset_dir alderley
+                     --phase train
+                     --experiment_name alderley_exp
+                     --batch_size 8 
+                     --load_size 286 
+                     --fine_size 256 
+                     --use_uncertainty True
+
+```
+
+```bash
+
+# BDD100k (256x512)
+python main_uncer.py --dataset_dir bdd100k 
+                     --phase train
+                     --experiment_name bdd_exp
+                     --batch_size 8 
+                     --load_size 286
+                     --fine_size 256 
+                     --use_uncertainty True
+
+```
+
+## Test
+
+```bash
+
+# Alderley (256x512)
+python main_uncer.py --dataset_dir alderley
+                     --phase test
+                     --experiment_name alderley_exp
+                     --batch_size 1 
+                     --load_size 286 
+                     --fine_size 256 
+                    
+```
+
+```bash
+
+# BDD100k (256x512)
+python main_uncer.py --dataset_dir bdd100k
+                     --phase test
+                     --experiment_name bdd_exp
+                     --batch_size 1 
+                     --load_size 286 
+                     --fine_size 256 
+                    
+
+```
+## Additional results
+<img src="./assets/augan_result.png" width="800">
+
+More results in [paper](https://www.bmvc2021-virtualconference.com/assets/papers/1443.pdf) and [supplementary]()
+
+## Uncertainty map 
+<img src="./assets/augan_uncer.png" width="800">
+
+## **Citation**
+If our code is helpful your research, please cite our paper:
+```
+@article{kwak2021adverse,
+  title={Adverse weather image translation with asymmetric and uncertainty-aware GAN},
+  author={Kwak, Jeong-gi and Jin, Youngsaeng and Li, Yuanming and Yoon, Dongsik and Kim, Donghyeon and Ko, Hanseok},
+  journal={arXiv preprint arXiv:2112.04283},
+  year={2021}
+}
+```
+## Acknowledgments
+Our code is bulided upon the [ForkGAN](https://www.ecva.net/papers/eccv_2020/papers_ECCV/papers/123480154.pdf) implementation.

cc.sh

@@ -0,0 +1,7 @@
+python main.py --dataset_dir swim \
+                     --phase test \
+                     --experiment_name bdd_exp \
+                     --batch_size 1  \
+                     --which_direction BtoA \
+                     --load_size 286 \
+                     --fine_size 256 

check.zip

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1fc7f6d4f5f9c503bc69e1fdb454d8cc4f652d8c8966875dcd46aab83cdd0ff4
+size 173513070

loss_utils.py

@@ -0,0 +1,51 @@
+import tensorflow as tf
+
+epsilon = 1e-7
+
+def conf_criterion_lp(im1, im2, conf_sigma):  # factorized laplacian distribution
+    loss = tf.abs(im1 - im2)
+    if conf_sigma is not None:
+        loss = loss * 2 / (conf_sigma + epsilon) + tf.log(conf_sigma * 2 + epsilon)
+        loss = tf.reduce_mean(loss)
+    else:
+        loss = tf.reduce_mean(loss)
+
+    return loss
+
+
+
+def conf_criterion(im1, im2, conf_sigma):  # gaussian distribution
+    loss = tf.abs(im1 - im2)
+    if conf_sigma is not None:
+        loss = tf.math.exp(-conf_sigma) * 5 * loss + conf_sigma / 2
+        loss = tf.reduce_mean(loss)
+    else:
+        loss = tf.reduce_mean(loss)
+
+    return loss
+
+
+def abs_criterion(in_, target):
+    return tf.reduce_mean(tf.abs(in_ - target))
+
+
+def mae_criterion(in_, target):
+    return tf.reduce_mean((in_ - target) ** 2)
+
+
+def sce_criterion(logits, labels):
+    return tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(logits=logits, labels=labels))
+
+
+def mae_criterion_list(in_, target):
+    loss = 0.0
+    for i in range(len(target)):
+        loss += tf.reduce_mean((in_[i] - target[i]) ** 2)
+    return loss / len(target)
+
+
+def sce_criterion_list(logits, labels):
+    loss = 0.0
+    for i in range(len(labels)):
+        loss += tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(logits=logits[i], labels=labels[i]))
+    return loss / len(labels)

main.py

@@ -0,0 +1,193 @@
+import argparse
+import tensorflow as tf
+import os
+from utils import *
+from AUGAN import AUGAN
+from ops import *
+import time
+
+parser = argparse.ArgumentParser(description="")
+parser.add_argument(
+    "--dataset_dir", dest="dataset_dir", default="bdd100k", help="path of the dataset"
+)
+parser.add_argument(
+    "--experiment_name",
+    dest="experiment_name",
+    type=str,
+    default="bdd_exp",
+    help="name of experiment",
+)
+parser.add_argument("--epoch", dest="epoch", type=int, default=20, help="# of epoch")
+parser.add_argument(
+    "--epoch_step",
+    dest="epoch_step",
+    type=int,
+    default=10,
+    help="# of epoch to decay lr",
+)
+parser.add_argument(
+    "--batch_size", dest="batch_size", type=int, default=1, help="# images in batch"
+)
+parser.add_argument(
+    "--train_size",
+    dest="train_size",
+    type=int,
+    default=1e8,
+    help="# images used to train",
+)
+parser.add_argument(
+    "--load_size",
+    dest="load_size",
+    type=int,
+    default=286,
+    help="scale images to this size",
+)
+parser.add_argument(
+    "--fine_size",
+    dest="fine_size",
+    type=int,
+    default=256,
+    help="then crop to this size",
+)
+parser.add_argument(
+    "--ngf",
+    dest="ngf",
+    type=int,
+    default=64,
+    help="# of gen filters in first conv layer",
+)
+parser.add_argument(
+    "--ndf",
+    dest="ndf",
+    type=int,
+    default=64,
+    help="# of discri filters in first conv layer",
+)
+parser.add_argument(
+    "--n_d", dest="n_d", type=int, default=2, help="# of discriminators"
+)
+parser.add_argument(
+    "--n_scale", dest="n_scale", type=int, default=2, help="# of scales"
+)
+parser.add_argument(
+    "--gpu", dest="gpu", type=int, default=0, help="# index of gpu device"
+)
+parser.add_argument(
+    "--input_nc", dest="input_nc", type=int, default=3, help="# of input image channels"
+)
+parser.add_argument(
+    "--output_nc",
+    dest="output_nc",
+    type=int,
+    default=3,
+    help="# of output image channels",
+)
+parser.add_argument(
+    "--lr", dest="lr", type=float, default=0.0002, help="initial learning rate for adam"
+)
+parser.add_argument(
+    "--beta1", dest="beta1", type=float, default=0.5, help="momentum term of adam"
+)
+parser.add_argument(
+    "--which_direction", dest="which_direction", default="AtoB", help="AtoB or BtoA "
+)
+parser.add_argument("--phase", dest="phase", default="test", help="train, test")
+parser.add_argument(
+    "--save_freq",
+    dest="save_freq",
+    type=int,
+    default=1000,
+    help="save a model every save_freq iterations",
+)
+parser.add_argument(
+    "--print_freq",
+    dest="print_freq",
+    type=int,
+    default=100,
+    help="print the debug information every print_freq iterations",
+)
+parser.add_argument(
+    "--L1_lambda",
+    dest="L1_lambda",
+    type=float,
+    default=10.0,
+    help="weight on L1 term in objective",
+)
+parser.add_argument(
+    "--conf_lambda",
+    dest="conf_lambda",
+    type=float,
+    default=1.0,
+    help="weight on L1 term in objective",
+)
+parser.add_argument(
+    "--use_resnet",
+    dest="use_resnet",
+    type=bool,
+    default=True,
+    help="generation network using reidule block",
+)
+parser.add_argument(
+    "--use_lsgan",
+    dest="use_lsgan",
+    type=bool,
+    default=True,
+    help="gan loss defined in lsgan",
+)
+parser.add_argument(
+    "--use_uncertainty",
+    dest="use_uncertainty",
+    type=bool,
+    default=True,
+    help="max size of image pool, 0 means do not use image pool",
+)
+parser.add_argument(
+    "--max_size",
+    dest="max_size",
+    type=int,
+    default=50,
+    help="max size of image pool, 0 means do not use image pool",
+)
+parser.add_argument(
+    "--continue_train",
+    dest="continue_train",
+    type=bool,
+    default=False,
+    help="if continue training, load the latest model: 1: true, 0: false",
+)
+parser.add_argument(
+    "--save_conf",
+    dest="save_conf",
+    type=bool,
+    default=False,
+    help="save conf map in test phase",
+)
+args = parser.parse_args()
+
+os.environ["CUDA_VISIBLE_DEVICES"] = str(args.gpu)
+
+
+def main(_):
+
+    set_path(args, args.experiment_name)
+
+    tfconfig = tf.compat.v1.ConfigProto(allow_soft_placement=True)
+    tfconfig.gpu_options.allow_growth = True
+    with tf.compat.v1.Session(config=tfconfig) as sess:
+        model = AUGAN(sess, args)
+        # show_all_variables()
+        # model.train(args) if args.phase == 'train' \
+        #     else model.test(args)
+
+        if args.phase == "train":
+            model.train(args)
+        elif args.phase == "test":
+            model.test(args)
+        elif args.phase == "convert":
+            model.convert_image(args, "inf_data/b1ca2e5d-84cf9134.jpg", "out")
+        else:
+            raise Exception("Give a phase")
+
+
+if __name__ == "__main__":
+    tf.compat.v1.app.run()

models.py

@@ -0,0 +1,178 @@
+from collections import namedtuple
+from utils import *
+from ops import *
+import time
+from glob import glob
+
+
+def gaussian_noise_layer(input_layer, std):
+    noise = tf.random.normal(
+        shape=tf.shape(input_layer), mean=0.0, stddev=std, dtype=tf.float32
+    )
+    return input_layer + noise
+
+
+def generator_resnet(image, options, transfer=False, reuse=False, name="generator"):
+    with tf.compat.v1.variable_scope(name):
+        if reuse:
+            tf.compat.v1.get_variable_scope().reuse_variables()
+        else:
+            assert tf.compat.v1.get_variable_scope().reuse is False
+
+        def residule_block_dilated(x, dim, ks=3, s=1, name="res", down=False):
+            if down:
+                dim = dim * 2
+            y = instance_norm(
+                dilated_conv2d(x, dim, ks, s, padding="SAME", name=name + "_c1"),
+                name + "_bn1",
+            )
+            y = tf.nn.relu(y)
+            y = instance_norm(
+                dilated_conv2d(y, dim, ks, s, padding="SAME", name=name + "_c2"),
+                name + "_bn2",
+            )
+            out = y + x
+            if down:
+                out = tf.nn.relu(
+                    instance_norm(
+                        conv2d(out, dim // 2, 3, 1, name=name + "_down_c"),
+                        name + "_in_down",
+                    )
+                )
+            return out
+
+        def residual_block(x_init, dim, ks=3, s=1, name="resblock", down=False):
+            with tf.compat.v1.variable_scope(name):
+                if down:
+                    dim = dim * 2
+
+                with tf.compat.v1.variable_scope("res1"):
+                    x = instance_norm(
+                        conv2d(x_init, dim, ks, s, padding="SAME", name=name + "_c1"),
+                        name + "_in1",
+                    )
+                    x = tf.nn.relu(x)
+
+                with tf.compat.v1.variable_scope("res2"):
+
+                    x = instance_norm(
+                        conv2d(x, dim, ks, s, padding="SAME", name=name + "_c2"),
+                        name + "_in2",
+                    )
+
+                out = x + x_init
+
+                if down:
+                    out = tf.nn.relu(
+                        instance_norm(
+                            conv2d(out, dim // 2, 3, 1, name=name + "_down_c"),
+                            name + "_in_down",
+                        )
+                    )
+                return out
+
+        ### Encoder architecture
+        c0 = tf.pad(image, [[0, 0], [3, 3], [3, 3], [0, 0]], "REFLECT")
+        c1 = tf.nn.relu(
+            instance_norm(
+                conv2d(c0, options.gf_dim, 7, 1, padding="VALID", name="g_e1_c"),
+                "g_e1_bn",
+            )
+        )
+        c2 = tf.nn.relu(
+            instance_norm(
+                conv2d(c1, options.gf_dim * 2, 3, 2, name="g_e2_c"), "g_e2_bn"
+            )
+        )
+        c3 = tf.nn.relu(
+            instance_norm(
+                conv2d(c2, options.gf_dim * 4, 3, 2, name="g_e3_c"), "g_e3_bn"
+            )
+        )
+        r1 = residule_block_dilated(c3, options.gf_dim * 4, name="g_r1")
+        r2 = residule_block_dilated(r1, options.gf_dim * 4, name="g_r2")
+        r3 = residule_block_dilated(r2, options.gf_dim * 4, name="g_r3")
+        r4 = residule_block_dilated(r3, options.gf_dim * 4, name="g_r4")
+        # r5 = residule_block_dilated(r4, options.gf_dim * 4, name='g_r5')
+
+        if transfer:
+            t1 = residual_block(r4, options.gf_dim * 4, name="g_t1")
+            t2 = residual_block(t1, options.gf_dim * 4, name="g_t2")
+            t3 = residual_block(t2, options.gf_dim * 4, name="g_t3")
+            t4 = residual_block(t3, options.gf_dim * 4, name="g_t4")
+            # feature = tf.concat([r4, t4], axis=3, name='g_concat')
+            # down = True
+            feature = t4
+        else:
+            feature = r4
+            t4 = None
+            down = False
+
+        ### translation decoder architecture
+        r6 = residule_block_dilated(feature, options.gf_dim * 4, name="g_r6")
+        r7 = residule_block_dilated(r6, options.gf_dim * 4, name="g_r7")
+        r8 = residule_block_dilated(r7, options.gf_dim * 4, name="g_r8")
+        r9 = residule_block_dilated(r8, options.gf_dim * 4, name="g_r9")
+        d1 = deconv2d(r9, options.gf_dim * 2, 3, 2, name="g_d1_dc")
+        d1 = tf.nn.relu(instance_norm(d1, "g_d1_bn"))
+        d2 = deconv2d(d1, options.gf_dim, 3, 2, name="g_d2_dc")
+        d2 = tf.nn.relu(instance_norm(d2, "g_d2_bn"))
+        d2 = tf.pad(d2, [[0, 0], [3, 3], [3, 3], [0, 0]], "REFLECT")
+        pred = tf.nn.tanh(
+            conv2d(d2, options.output_c_dim, 7, 1, padding="VALID", name="g_pred_c")
+        )
+
+        ### reconstruction decoder architecture
+        r5 = gaussian_noise_layer(r4, 0.02)
+        r6_rec = residule_block_dilated(r5, options.gf_dim * 4, name="g_r6_rec")
+        r6_rec = gaussian_noise_layer(r6_rec, 0.02)
+        r7_rec = residule_block_dilated(r6_rec, options.gf_dim * 4, name="g_r7_rec")
+        r8_rec = residule_block_dilated(r7_rec, options.gf_dim * 4, name="g_r8_rec")
+        r9_rec = residule_block_dilated(r8_rec, options.gf_dim * 4, name="g_r9_rec")
+        d1_rec = deconv2d(r9_rec, options.gf_dim * 2, 3, 2, name="g_d1_dc_rec")
+        d1_rec = tf.nn.relu(instance_norm(d1_rec, "g_d1_bn_rec"))
+        d2_rec = deconv2d(d1_rec, options.gf_dim, 3, 2, name="g_d2_dc_rec")
+        d2_rec = tf.nn.relu(instance_norm(d2_rec, "g_d2_bn_rec"))
+        d2_rec = tf.pad(d2_rec, [[0, 0], [3, 3], [3, 3], [0, 0]], "REFLECT")
+        pred_rec = tf.nn.tanh(
+            conv2d(
+                d2_rec, options.output_c_dim, 7, 1, padding="VALID", name="g_pred_c_rec"
+            )
+        )
+
+        ## confidence prediction
+
+        if transfer:
+
+            d_conf = deconv2d(d1, options.gf_dim, 3, 2, name="g_d_dc_conf")
+            d_conf = tf.nn.relu(instance_norm(d_conf, "g_d_bn_conf"))
+            d_conf = tf.pad(d_conf, [[0, 0], [3, 3], [3, 3], [0, 0]], "REFLECT")
+            pred_conf = tf.nn.softplus(
+                conv2d(d_conf, 1, 7, 1, padding="VALID", name="g_pred_c_conf")
+            )
+
+        else:
+            pred_conf = None
+
+        return pred, pred_rec, r4, t4, pred_conf
+
+
+def discriminator(image, options, n_scale=2, reuse=False, name="discriminator"):
+    images = []
+    for i in range(n_scale):
+        images.append(
+            tf.compat.v1.image.resize_bicubic(
+                image, [get_shape(image)[1] // (2**i), get_shape(image)[2] // (2**i)]
+            )
+        )
+    with tf.compat.v1.variable_scope(name):
+        if reuse:
+            tf.compat.v1.get_variable_scope().reuse_variables()
+        else:
+            assert tf.compat.v1.get_variable_scope().reuse is False
+        images = dis_down(images, 4, 2, n_scale, options.df_dim, "d_h0_conv_scale_")
+        images = dis_down(images, 4, 2, n_scale, options.df_dim * 2, "d_h1_conv_scale_")
+        images = dis_down(images, 4, 2, n_scale, options.df_dim * 4, "d_h2_conv_scale_")
+        images = dis_down(images, 4, 2, n_scale, options.df_dim * 8, "d_h3_conv_scale_")
+        images = final_conv(images, n_scale, "d_pred_scale_")
+        return images

ops.py

@@ -0,0 +1,246 @@
+import tensorflow as tf
+
+# import tensorflow.contrib.slim as slim
+import tf_slim as slim
+import math
+import pprint
+
+pp = pprint.PrettyPrinter()
+get_stddev = lambda x, k_h, k_w: 1 / math.sqrt(k_w * k_h * x.get_shape()[-1])
+# import tensorflow.contrib as tf_contrib
+
+# weight_init = tf_contrib.layers.xavier_initializer()
+weight_init = tf.initializers.GlorotUniform()
+weight_regularizer = None
+
+
+def batch_norm(x, name="batch_norm"):
+    # return tf.contrib.layers.batch_norm(
+    #     x, decay=0.9, updates_collections=None, epsilon=1e-5, scale=True, scope=name
+    # )
+    return tf.keras.layers.BatchNormalization(
+        momentum=0.9, epsilon=1e-5, scale=True, name=name
+    )(x)
+
+
+def instance_norm(input, name="instance_norm"):
+    with tf.compat.v1.variable_scope(name):
+        depth = input.get_shape()[3]
+        scale = tf.compat.v1.get_variable(
+            "scale",
+            [depth],
+            initializer=tf.keras.initializers.RandomNormal(
+                mean=1.0, stddev=0.02, seed=None
+            ),
+        )
+        offset = tf.compat.v1.get_variable(
+            "offset", [depth], initializer=tf.constant_initializer(0.0)
+        )
+        mean, variance = tf.nn.moments(input, axes=[1, 2], keepdims=True)
+        epsilon = 1e-5
+        inv = tf.math.rsqrt(variance + epsilon)
+        normalized = (input - mean) * inv
+        return scale * normalized + offset
+
+
+def conv2d(input_, output_dim, ks=4, s=2, stddev=0.02, padding="SAME", name="conv2d"):
+    with tf.compat.v1.variable_scope(name):
+        return slim.conv2d(
+            input_,
+            output_dim,
+            ks,
+            s,
+            padding=padding,
+            activation_fn=None,
+            weights_initializer=tf.keras.initializers.TruncatedNormal(stddev=stddev),
+            biases_initializer=None,
+        )
+
+
+def deconv2d(input_, output_dim, ks=4, s=2, stddev=0.02, name="deconv2d"):
+    with tf.compat.v1.variable_scope(name):
+        return slim.conv2d_transpose(
+            input_,
+            output_dim,
+            ks,
+            s,
+            padding="SAME",
+            activation_fn=None,
+            weights_initializer=tf.keras.initializers.TruncatedNormal(stddev=stddev),
+            biases_initializer=None,
+        )
+
+
+def dilated_conv2d(
+    input_, output_dim, ks=3, s=2, stddev=0.02, padding="SAME", name="conv2d"
+):
+    with tf.compat.v1.variable_scope(name):
+        batch, in_height, in_width, in_channels = [int(d) for d in input_.get_shape()]
+        filter = tf.compat.v1.get_variable(
+            "filter",
+            [ks, ks, in_channels, output_dim],
+            dtype=tf.float32,
+            initializer=tf.random_normal_initializer(0, stddev),
+        )
+        conv = tf.nn.atrous_conv2d(input_, filter, rate=s, padding=padding, name=name)
+
+        return conv
+
+
+def one_step(x, ch, kernel, stride, name):
+    return lrelu(
+        instance_norm(
+            conv2d(x, ch, kernel, stride, name=name + "_first_c"), name + "_first_bn"
+        )
+    )
+
+
+def one_step_dilated(x, ch, kernel, stride, name):
+    return lrelu(
+        instance_norm(
+            dilated_conv2d(x, ch, kernel, stride, name=name + "_first_c"),
+            name + "_first_bn",
+        )
+    )
+
+
+def num_steps(x, ch, kernel, stride, num_steps, name):
+    for i in range(num_steps):
+        x = lrelu(
+            instance_norm(
+                conv2d(x, ch, kernel, stride, name=name + "_c_" + str(i)),
+                name + "_bn_" + str(i),
+            )
+        )
+    return x
+
+
+def one_step_noins(x, ch, kernel, stride, name):
+    return lrelu(conv2d(x, ch, kernel, stride, name=name + "_first_c"))
+
+
+def num_steps_noins(x, ch, kernel, stride, num_steps, name):
+
+    for i in range(num_steps):
+        x = lrelu(conv2d(x, ch, kernel, stride, name=name + "_c_" + str(i)))
+    return x
+
+
+def dis_down(images, kernel_size, stride, n_scale, ch, name):
+    backpack = images[0]
+    for i in range(n_scale):
+        if i == n_scale - 1:
+            images[i] = num_steps(
+                backpack, ch, kernel_size, stride, n_scale, name + str(i)
+            )
+        else:
+            images[i] = one_step_dilated(
+                images[i + 1], ch, kernel_size, 1, name + str(i)
+            )
+    return images
+
+
+def dis_down_noins(images, kernel_size, stride, n_scale, ch, name):
+    backpack = images[0]
+    for i in range(n_scale):
+        if i == n_scale - 1:
+            images[i] = num_steps_noins(
+                backpack, ch, kernel_size, stride, n_scale, name + str(i)
+            )
+        else:
+            images[i] = one_step_noins(images[i + 1], ch, kernel_size, 1, name + str(i))
+    return images
+
+
+def final_conv(images, n_scale, name):
+    for i in range(n_scale):
+        images[i] = conv2d(images[i], 1, s=1, name=name + str(i))
+    return images
+
+
+def lrelu(x, leak=0.2, name="lrelu"):
+    return tf.maximum(x, leak * x)
+
+
+def linear(input_, output_size, scope=None, stddev=0.02, bias_start=0.0, with_w=False):
+    with tf.compat.v1.variable_scope(scope or "Linear"):
+        matrix = tf.get_variable(
+            "Matrix",
+            [input_.get_shape()[-1], output_size],
+            tf.float32,
+            tf.random_normal_initializer(stddev=stddev),
+        )
+        bias = tf.get_variable(
+            "bias", [output_size], initializer=tf.constant_initializer(bias_start)
+        )
+        if with_w:
+            return tf.matmul(input_, matrix) + bias, matrix, bias
+        else:
+            return tf.matmul(input_, matrix) + bias
+
+
+def get_ones_like(logit):
+    target = []
+    for i in range(len(logit)):
+        target.append(tf.ones_like(logit[i]))
+    return target
+
+
+def get_zeros_like(logit):
+    target = []
+    for i in range(len(logit)):
+        target.append(tf.zeros_like(logit[i]))
+    return target
+
+
+def conv(
+    x,
+    channels,
+    kernel=4,
+    stride=2,
+    pad=0,
+    pad_type="zero",
+    use_bias=True,
+    scope="conv_0",
+):
+    with tf.compat.v1.variable_scope(scope):
+        if pad_type == "zero":
+            x = tf.pad(x, [[0, 0], [pad, pad], [pad, pad], [0, 0]])
+        if pad_type == "reflect":
+            x = tf.pad(x, [[0, 0], [pad, pad], [pad, pad], [0, 0]], mode="REFLECT")
+
+        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 reduce_sum(input_tensor, axis=None, keepdims=False):
+    try:
+        return tf.reduce_sum(input_tensor, axis=axis, keepdims=keepdims)
+    except:
+        return tf.reduce_sum(input_tensor, axis=axis, keep_dims=keepdims)
+
+
+def get_shape(inputs, name=None):
+    name = "shape" if name is None else name
+    with tf.name_scope(name):
+        static_shape = inputs.get_shape().as_list()
+        dynamic_shape = tf.shape(inputs)
+        shape = []
+        for i, dim in enumerate(static_shape):
+            dim = dim if dim is not None else dynamic_shape[i]
+            shape.append(dim)
+        return shape
+
+
+def show_all_variables():
+    model_vars = tf.trainable_variables()
+    slim.model_analyzer.analyze_vars(model_vars, print_info=True)

parser.py

@@ -0,0 +1,18 @@
+import json
+import shutil
+
+
+with open('C:/jg/github_code/ForkGAN/bdd100k/labels/bdd100k_labels_images_train.json') as json_file:
+    json_data = json.load(json_file)
+
+for item in json_data:
+    item_path = 'C:/jg/github_code/ForkGAN/bdd100k/images/100k/train/'+ item['name']
+    print(item['name'])
+    if item['attributes']['timeofday'] == 'daytime':
+        shutil.copy(item_path, 'C:/jg/github_code/ForkGAN/bdd100k/images/daytime/'+item['name'])
+
+    elif item['attributes']['timeofday'] == 'night':
+        shutil.copy(item_path, 'C:/jg/github_code/ForkGAN/bdd100k/images/night/'+item['name'])
+
+    else :
+        shutil.copy(item_path, 'C:/jg/github_code/ForkGAN/bdd100k/images/else/' + item['name'])

requirements.txt

@@ -0,0 +1,4 @@
+pillow==6.0.0
+scipy==1.1.0
+numpy
+matplotlib              

utils.py

@@ -0,0 +1,182 @@
+# import scipy.misc
+from PIL import Image
+import numpy as np
+import copy
+import os
+
+
+class ImagePool(object):
+    def __init__(self, maxsize=50):
+        self.maxsize = maxsize
+        self.num_img = 0
+        self.images = []
+
+    def __call__(self, image):
+        if self.maxsize <= 0:
+            return image
+        if self.num_img < self.maxsize:
+            self.images.append(image)
+            self.num_img += 1
+            return image
+        if np.random.rand() > 0.5:
+            idx = int(np.random.rand() * self.maxsize)
+            tmp1 = copy.copy(self.images[idx])[0]
+            self.images[idx][0] = image[0]
+            idx = int(np.random.rand() * self.maxsize)
+            tmp2 = copy.copy(self.images[idx])[1]
+            self.images[idx][1] = image[1]
+            return [tmp1, tmp2]
+        else:
+            return image
+
+
+def load_test_data(image_path, fine_size=256):
+    img = Image.open(image_path)
+    img = img.resize((fine_size * 2, fine_size))
+    img = np.array(img)
+    # Normalize image to the range [-1, 1]
+    img = img / 127.5 - 1
+
+    return img
+
+
+def check_folder(path):
+    if not os.path.exists(path):
+        os.mkdir(path)
+
+
+def load_train_data(image_path, load_size=286, fine_size=256, is_testing=False):
+    img_A = Image.open(image_path[0])
+    img_B = Image.open(image_path[1])
+
+    if not is_testing:
+        # Resize images using PIL
+        img_A = img_A.resize((load_size * 2, load_size))
+        img_B = img_B.resize((load_size * 2, load_size))
+
+        # Random crop
+        h1 = int(np.ceil(np.random.uniform(1e-2, load_size - fine_size)))
+        w1 = int(np.ceil(np.random.uniform(1e-2, (load_size - fine_size) * 2)))
+        img_A = np.array(img_A.crop((w1, h1, w1 + fine_size * 2, h1 + fine_size)))
+        img_B = np.array(img_B.crop((w1, h1, w1 + fine_size * 2, h1 + fine_size)))
+
+        # Random horizontal flip
+        if np.random.random() > 0.5:
+            img_A = np.fliplr(img_A)
+            img_B = np.fliplr(img_B)
+    else:
+        # Resize images using PIL for testing
+        img_A = img_A.resize((fine_size * 2, fine_size))
+        img_B = img_B.resize((fine_size * 2, fine_size))
+
+    # Normalize images to the range [-1, 1]
+    img_A = img_A / 127.5 - 1.0
+    img_B = img_B / 127.5 - 1.0
+
+    # Concatenate images along the channel axis
+    img_AB = np.concatenate((img_A, img_B), axis=2)
+
+    return img_AB
+
+
+# -----------------------------
+
+
+def get_image(image_path, image_size, is_crop=True, resize_w=64, is_grayscale=False):
+    return transform(
+        load_image(image_path, is_grayscale), image_size, is_crop, resize_w
+    )
+
+
+def save_images(images, size, image_path):
+    return imsave(images, size, image_path)
+
+
+def load_image(path, is_grayscale=False):
+    if is_grayscale:
+        return np.array(Image.open(path).convert("L")).astype(np.float)
+    else:
+        return np.array(Image.open(path).convert("RGB")).astype(np.float)
+
+
+def merge_images(images, size):
+    return inverse_transform(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[j * h : j * h + h, i * w : i * w + w, :] = image
+
+    return img
+
+
+def imsave(image, size, path):
+    # Convert images to uint8 format and adjust the range
+    image = ((image + 1.0) * 127.5).astype(np.uint8)
+
+    # Merge images
+    # merged_image = merge(images, size).astype(np.uint8)
+
+    # Create a PIL Image from the numpy array
+    pil_image = Image.fromarray(image)
+
+    # Save the image using PIL
+    pil_image.save(path)
+
+    return None
+
+
+def center_crop(x, crop_h, crop_w, resize_h=64, resize_w=64):
+    if crop_w is None:
+        crop_w = crop_h
+    h, w = x.shape[:2]
+    j = int(round((h - crop_h) / 2.0))
+    i = int(round((w - crop_w) / 2.0))
+
+    # Use PIL for resizing
+    cropped_image = Image.fromarray(x[j : j + crop_h, i : i + crop_w].astype(np.uint8))
+    cropped_image = cropped_image.resize((resize_w, resize_h))
+
+    return np.array(cropped_image) / 127.5 - 1.0
+
+
+def transform(image, npx=64, is_crop=True, resize_w=64):
+    # npx: # of pixels width/height of image
+    if is_crop:
+        cropped_image = center_crop(image, npx, resize_w=resize_w)
+    else:
+        cropped_image = image
+    return np.array(cropped_image) / 127.5 - 1.0
+
+
+def inverse_transform(images):
+    return (images + 1.0) / 2.0
+
+
+def norm_img(img):
+    img = img / np.linalg.norm(img)
+    img = (img * 2.0) - 1.0
+
+    return img
+
+
+def set_path(args, experiment_name):
+    args.checkpoint_dir = f"./check/{experiment_name}"
+    args.sample_dir = f"./check/{experiment_name}/sample"
+    if args.which_direction == "AtoB":
+        args.test_dir = f"./check/{experiment_name}/testa2b"
+    else:
+        args.test_dir = f"./check/{experiment_name}/testb2a"
+    args.conf_dir = f"./check/{experiment_name}/conf"
+    if not os.path.exists(args.checkpoint_dir):
+        os.makedirs(args.checkpoint_dir)
+    if not os.path.exists(args.sample_dir):
+        os.makedirs(args.sample_dir)
+    if not os.path.exists(args.test_dir):
+        os.makedirs(args.test_dir)
+    if not os.path.exists(args.conf_dir):
+        os.makedirs(args.conf_dir)