import os
from time import time
import pickle

import cv2
import numpy as np
from PIL import Image
import scipy.misc
import json
import sys
import tensorflow as tf
from tensorflow.python.client import timeline

from graph_single import build_multi_tower_graph, build_single_graph
from input_pipeline import build_input_queue_paired, build_input_queue_paired_test, resize_and_padding_mask_image, \
    thicken_drawings
from text_processing import preprocess_sentence, load_vocab_dict_from_file
from config import Config

tf.logging.set_verbosity(tf.logging.INFO)
inception_v4_ckpt_path = 'model/inception-cartoon'


def log(name, arr):
    print(name, ', ', arr.shape, ', max:', np.max(arr), ', min:', np.min(arr))


def print_parameter_count(verbose=False):
    total_parameters = 0
    for variable in tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, scope='generator'):
        # shape is an array of tf.Dimension
        shape = variable.get_shape()
        # print(len(shape))
        variable_parametes = 1
        for dim in shape:
            # print(dim)
            variable_parametes *= dim.value
        if verbose and len(shape) > 1:
            print(shape)
            print(variable_parametes)
        total_parameters += variable_parametes
    print('generator')
    print('total_parameters', total_parameters)

    total_parameters = 0
    for variable in tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, scope='discriminator'):
        # shape is an array of tf.Dimension
        shape = variable.get_shape()
        # print(len(shape))
        variable_parametes = 1
        for dim in shape:
            # print(dim)
            variable_parametes *= dim.value
        if verbose and len(shape) > 1:
            print(shape)
            print(variable_parametes)
        total_parameters += variable_parametes
    print('discriminator')
    print('total_parameters', total_parameters)


def train(**kwargs):
    status = 0

    # Roll out the parameters
    appendix = Config.resume_from
    batch_size = Config.batch_size
    max_iter_step = Config.max_iter_step
    Diters = Config.disc_iterations
    ld = Config.ld
    optimizer = Config.optimizer
    lr_G = Config.lr_G
    lr_D = Config.lr_D
    num_gpu = Config.num_gpu
    log_dir = Config.log_dir
    ckpt_dir = Config.ckpt_dir
    data_format = Config.data_format
    distance_map = Config.distance_map
    small_img = Config.small_img
    LSTM_hybrid = Config.LSTM_hybrid
    block_type = Config.block_type
    summary_write_freq = Config.summary_write_freq
    save_model_freq = Config.save_model_freq
    count_left_time_freq = Config.count_left_time_freq
    # count_inception_score_freq = Config.count_inception_score_freq
    vocab_size = Config.vocab_size

    distance_map = distance_map != 0
    small = small_img != 0
    LSTM_hybrid = LSTM_hybrid != 0
    batch_portion = np.array([1, 1, 1, 1], dtype=np.int32)

    iter_from = kwargs['iter_from']

    # Time counter
    prev_time = float("-inf")
    curr_time = float("-inf")

    tf.reset_default_graph()
    print('Iteration starts from: %d' % iter_from)

    # assert inception_score.softmax.graph != tf.get_default_graph()
    # inception_score._init_inception()

    counter = tf.Variable(initial_value=iter_from, dtype=tf.int32, trainable=False)
    counter_addition_op = tf.assign_add(counter, 1, use_locking=True)

    # Construct data queue
    with tf.device('/cpu:0'):
        images, sketches, image_paired_class_ids, _, _, _, text_vocab_indiceses = build_input_queue_paired(
            mode='train',
            batch_size=batch_size * num_gpu,
            data_format=data_format,
            distance_map=distance_map,
            small=small, capacity=2 ** 11)  # images/sketches [2, 3, H, W], text_vocab_indiceses [2, 15]
    with tf.device('/cpu:0'):
        images_d, _, image_paired_class_ids_d, _, _, _, _ = build_input_queue_paired(
            mode='train',
            batch_size=batch_size * num_gpu,
            data_format=data_format,
            distance_map=distance_map,
            small=small, capacity=2 ** 11)  # [2, 3, H, W]

    opt_g, opt_d, loss_g, loss_d, merged_all = build_multi_tower_graph(
        images, sketches, images_d,
        image_paired_class_ids, image_paired_class_ids_d,
        text_vocab_indiceses,
        LSTM_hybrid=LSTM_hybrid,
        vocab_size=vocab_size,
        batch_size=batch_size, num_gpu=num_gpu, batch_portion=batch_portion, training=True,
        learning_rates={
            "generator": lr_G,
            "discriminator": lr_D,
        },
        counter=counter, max_iter_step=max_iter_step,
        ld=ld, data_format=data_format,
        distance_map=distance_map,
        optimizer=optimizer,
        block_type=block_type)

    saver = tf.train.Saver(max_to_keep=100)
    # try:
    #     inception_loader = tf.train.Saver(tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope='InceptionV4'))
    #     perceptual_model_checkpoint_path = inception_v4_ckpt_path
    #     perceptual_model_path = tf.train.latest_checkpoint(perceptual_model_checkpoint_path)
    # except:
    #     inception_loader = None

    config = tf.ConfigProto(allow_soft_placement=True,
                            intra_op_parallelism_threads=4, inter_op_parallelism_threads=4)
    config.gpu_options.allow_growth = True

    with tf.Session(config=config) as sess:
        sess.run(tf.global_variables_initializer())
        sess.run(tf.local_variables_initializer())

        # if inception_loader is not None:
        #     print('Restore:', perceptual_model_path)
        #     inception_loader.restore(sess, perceptual_model_path)

        summary_writer = tf.summary.FileWriter(log_dir, sess.graph)
        if iter_from > 0:
            snapshot_loader = tf.train.Saver()
            print('Restore:', tf.train.latest_checkpoint(ckpt_dir))
            snapshot_loader.restore(sess, tf.train.latest_checkpoint(ckpt_dir))
            summary_writer.reopen()

        run_options = tf.RunOptions(trace_level=tf.RunOptions.NO_TRACE)
        run_metadata = tf.RunMetadata()

        print_parameter_count(verbose=False)

        coord = tf.train.Coordinator()
        threads = tf.train.start_queue_runners(sess=sess, coord=coord)

        sess.run([counter.assign(iter_from)])

        for i in range(iter_from, max_iter_step):
            if status == -1:
                break

            ## count left time
            if i % count_left_time_freq == 0:
                curr_time = time()
                elapsed = curr_time - prev_time
                print("Now at iteration %d. Elapsed time: %.5fs. Average time: %.5fs/iter"
                      % (i, elapsed, elapsed / 100.))

                if elapsed != float("inf"):
                    left_iter = max_iter_step - i
                    left_sec = left_iter * (elapsed / 100.)
                    left_day = int(left_sec / 24 / 60 / 60)
                    left_hour = int((left_sec - (24 * 60 * 60) * left_day) / 60 / 60)
                    left_min = int((left_sec - (24 * 60 * 60) * left_day - (60 * 60) * left_hour) / 60)
                    print("Left time:%dd %dh %dm" % (left_day, left_hour, left_min))

                prev_time = curr_time

            diters = Diters

            # Train Discriminator
            for j in range(diters):
                ## summary
                if i % summary_write_freq == 0 and j == 0:
                    _, merged, loss_d_out = sess.run([opt_d, merged_all, loss_d],
                                                     options=run_options,
                                                     run_metadata=run_metadata)
                    summary_writer.add_summary(merged, i)
                    summary_writer.flush()
                else:
                    _, loss_d_out = sess.run([opt_d, loss_d])
                # print('loss_d', loss_d_out)
                if np.isnan(np.sum(loss_d_out)):
                    status = -1
                    print("NaN occurred during training D")
                    return status

            # Train Generator
            if i % summary_write_freq == 0:
                _, merged, loss_g_out, counter_out, _ = sess.run(
                    [opt_g, merged_all, loss_g, counter, counter_addition_op],
                    options=run_options,
                    run_metadata=run_metadata)
                summary_writer.add_summary(merged, i)
                summary_writer.flush()
            else:
                _, loss_g_out, counter_out, _ = sess.run([opt_g, loss_g, counter, counter_addition_op])
            # print('loss_g', loss_g_out)
            if np.isnan(np.sum(loss_g_out)):
                status = -1
                print("NaN occurred during training G")
                return status

            ## save model
            if i % save_model_freq == save_model_freq - 1:
                saver.save(sess, os.path.join(ckpt_dir, 'model_{}.ckpt'.format(i)), global_step=i)
                print('Save model_{}.ckpt'.format(i))

        coord.request_stop()
        coord.join(threads)

    return status


def validation(**kwargs):
    # Roll out the parameters
    dataset_type = Config.dataset_type
    batch_size = Config.batch_size
    ckpt_dir = Config.ckpt_dir
    results_dir = Config.results_dir
    data_format = Config.data_format
    distance_map = Config.distance_map
    small_img = Config.small_img
    LSTM_hybrid = Config.LSTM_hybrid
    block_type = Config.block_type
    vocab_size = Config.vocab_size

    channel = 3
    distance_map = distance_map != 0
    small = small_img != 0
    LSTM_hybrid = LSTM_hybrid != 0

    if LSTM_hybrid:
        output_folder = os.path.join(results_dir, 'with_text')
    else:
        output_folder = os.path.join(results_dir, 'without_text')
    print(output_folder)
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)

    # Time counter
    prev_time = float("-inf")
    curr_time = float("-inf")
    # Construct data queue
    with tf.device('/cpu:0'):
        images, sketches, class_ids, categories, image_names, color_texts, text_vocab_indiceses \
            = build_input_queue_paired_test(
            mode=dataset_type,
            batch_size=batch_size, data_format=data_format,
            distance_map=distance_map, small=small, capacity=512)  # [2, 3, H, W]

    ret_list = build_single_graph(images, sketches, None,
                                  class_ids, None,
                                  text_vocab_indiceses,
                                  batch_size=batch_size, training=False,
                                  LSTM_hybrid=LSTM_hybrid,
                                  vocab_size=vocab_size,
                                  data_format=data_format,
                                  distance_map=distance_map,
                                  block_type=block_type)  # [image_gens, images, sketches]

    snapshot_loader = tf.train.Saver()

    tf_config = tf.ConfigProto(allow_soft_placement=True)
    tf_config.gpu_options.allow_growth = True
    with tf.Session(config=tf_config) as sess:
        sess.run(tf.global_variables_initializer())
        sess.run(tf.local_variables_initializer())

        print('Restore trained model:', tf.train.latest_checkpoint(ckpt_dir))
        snapshot_loader.restore(sess, tf.train.latest_checkpoint(ckpt_dir))
        counter = 0

        coord = tf.train.Coordinator()
        threads = tf.train.start_queue_runners(sess=sess, coord=coord)

        while True:
            try:
                generated_img, gt_image, input_sketch, category, image_name = sess.run(
                    [ret_list[0], ret_list[1], ret_list[2], categories, image_names])
            except Exception as e:
                print(e.args)
                break

            if counter % 100 == 0:
                curr_time = time()
                elapsed = curr_time - prev_time
                print(
                    "Now at iteration %d. Elapsed time: %.5fs." % (counter, elapsed))
                prev_time = curr_time

            if data_format == 'NCHW':
                generated_img = np.transpose(generated_img, (0, 2, 3, 1))
                gt_image = np.transpose(gt_image, (0, 2, 3, 1))
                input_sketch = np.transpose(input_sketch, (0, 2, 3, 1))

            # log('before, generated_img', generated_img)
            # log('before, gt_image', gt_image)
            # log('before, input_sketch', input_sketch)
            generated_img = ((generated_img + 1) / 2.) * 255
            gt_image = ((gt_image + 1) / 2.) * 255
            input_sketch = ((input_sketch + 1) / 2.) * 255
            generated_img = generated_img[:, :, :, ::-1].astype(np.uint8)
            gt_image = gt_image[:, :, :, ::-1].astype(np.uint8)
            input_sketch = input_sketch.astype(np.uint8)
            # log('after, generated_img', generated_img)
            # log('after, gt_image', gt_image)  # (2, H, W, 3)
            # log('after, input_sketch', input_sketch)

            for i in range(batch_size):
                this_prefix = '%s' % (category[i].decode('ascii'))
                img_out_filename = this_prefix + '_' + image_name[i].decode()[:-4] + '_output.png'
                img_gt_filename = this_prefix + '_' + image_name[i].decode()[:-4] + '_target.png'
                sketch_in_filename = this_prefix + '_' + image_name[i].decode()[:-4] + '_input.png'

                # Save file
                # file_path = os.path.join(output_folder, 'output_%d.jpg' % int(counter / batch_size))
                cv2.imwrite(os.path.join(output_folder, img_out_filename), generated_img[i])
                cv2.imwrite(os.path.join(output_folder, img_gt_filename), gt_image[i])
                cv2.imwrite(os.path.join(output_folder, sketch_in_filename), input_sketch[i])
                # output_img = np.zeros((img_dim * 2, img_dim * batch_size, channel))

                print('Saved file %s' % this_prefix)

            counter += 1

        coord.request_stop()
        coord.join(threads)


def test():
    SIZE = {True: (64, 64),
            False: (192, 192)}
    T = 15  # the longest length of text
    vocab_file = 'data/vocab.txt'
    test_data_base_dir = 'data'

    captions_base_dir = os.path.join(test_data_base_dir, 'captions')
    images_base_dir = os.path.join(test_data_base_dir, 'images')
    categories = os.listdir(captions_base_dir)
    categories.sort()
    print(categories)

    # Roll out the parameters
    batch_size = 1
    ckpt_dir = Config.ckpt_dir
    results_dir = Config.results_dir
    data_format = Config.data_format
    distance_map = Config.distance_map
    small_img = Config.small_img
    LSTM_hybrid = Config.LSTM_hybrid
    block_type = Config.block_type
    vocab_size = Config.vocab_size

    distance_map = distance_map != 0
    small = small_img != 0
    LSTM_hybrid = LSTM_hybrid != 0

    img_dim = SIZE[small]

    output_folder = results_dir
    os.makedirs(output_folder, exist_ok=True)

    vocab_dict = load_vocab_dict_from_file(vocab_file)

    input_images = tf.placeholder(tf.float32, shape=[1, 3, img_dim[0], img_dim[1]])  # [1, 3, H, W]
    class_ids = tf.placeholder(tf.int32, shape=(1,))  # (1, )
    text_vocab_indiceses = tf.placeholder(tf.int32, shape=[1, 15])  # [1, 15]

    ret_list = build_single_graph(input_images, input_images, None,
                                  class_ids, None,
                                  text_vocab_indiceses,
                                  batch_size=batch_size, training=False,
                                  LSTM_hybrid=LSTM_hybrid,
                                  vocab_size=vocab_size,
                                  data_format=data_format,
                                  distance_map=distance_map,
                                  block_type=block_type)  # [image_gens, images, sketches]

    snapshot_loader = tf.train.Saver()

    tf_config = tf.ConfigProto(allow_soft_placement=True)
    tf_config.gpu_options.allow_growth = True
    with tf.Session(config=tf_config) as sess:
        sess.run(tf.global_variables_initializer())
        sess.run(tf.local_variables_initializer())

        print('Restore trained model:', tf.train.latest_checkpoint(ckpt_dir))
        snapshot_loader.restore(sess, tf.train.latest_checkpoint(ckpt_dir))

        for cate in categories:
            testing_json = os.path.join(captions_base_dir, cate, 'test.json')
            fp = open(testing_json, "r")
            json_data = fp.read()
            json_data = json.loads(json_data)
            print(len(json_data), 'inference datas')

            for i in range(len(json_data)):
                input_name = json_data[i]['key']  # e.g. '228_1.png'
                input_text = json_data[i]['color_text']  # e.g. 'A yellow bus with blue window'

                sketch_path = os.path.join(images_base_dir, cate, 'sketch', input_name)
                sketch_image = Image.open(sketch_path)
                sketch_image = sketch_image.convert("RGB")
                # Resize
                if sketch_image.width != img_dim[0] or sketch_image.height != img_dim[1]:
                    margin_size = 0 if cate in ['road'] else 10
                    sketch_image = resize_and_padding_mask_image(sketch_image, img_dim[0],
                                                                 margin_size=margin_size).astype(np.float32)
                else:
                    sketch_image = np.array(sketch_image, dtype=np.float32)  # shape = [H, W, 3]

                if cate in ['house', 'road']:
                    sketch_image = thicken_drawings(sketch_image).astype(np.float32)  # shape = [H, W, 3]

                # Normalization
                sketch_image = sketch_image / 255.
                sketch_image = sketch_image * 2. - 1

                sketch_image = np.expand_dims(sketch_image, axis=0)  # shape = [1, H, W, 3]
                sketch_image = np.transpose(sketch_image, [0, 3, 1, 2])  # shape = [1, 3, H, W]

                class_id = categories.index(cate)
                class_id = np.array([class_id])

                vocab_indices = preprocess_sentence(input_text, vocab_dict, T)  # list
                vocab_indices = np.array(vocab_indices, dtype=np.int32)
                vocab_indices = np.expand_dims(vocab_indices, axis=0)  # shape = [1, 15]

                try:
                    # print('class_id', class_id)
                    # print('vocab_indices', vocab_indices)
                    generated_img, _, input_sketch = sess.run(
                        [ret_list[0], ret_list[1], ret_list[2]],
                        feed_dict={input_images: sketch_image,
                                   class_ids: class_id,
                                   text_vocab_indiceses: vocab_indices})
                except Exception as e:
                    print(e.args)
                    break

                if data_format == 'NCHW':
                    generated_img = np.transpose(generated_img, (0, 2, 3, 1))
                    input_sketch = np.transpose(input_sketch, (0, 2, 3, 1))

                # log('before, generated_img', generated_img)
                # log('before, input_sketch', input_sketch)
                generated_img = ((generated_img + 1) / 2.) * 255
                input_sketch = ((input_sketch + 1) / 2.) * 255
                generated_img = generated_img[:, :, :, ::-1].astype(np.uint8)
                input_sketch = input_sketch.astype(np.uint8)
                # log('after, generated_img', generated_img)
                # log('after, input_sketch', input_sketch)

                img_out_filename = cate + '_' + input_name[:-4] + '_output.png'
                sketch_in_filename = cate + '_' + input_name[:-4] + '_input.png'

                # Save file
                cv2.imwrite(os.path.join(output_folder, img_out_filename), generated_img[0])
                cv2.imwrite(os.path.join(output_folder, sketch_in_filename), input_sketch[0])

                print('Saved file %s' % img_out_filename)


def inference(img_name, instruction):
    wild_data_base_dir = 'examples'
    wild_text = instruction

    wild_cate = img_name[:img_name.find('.png')]

    SIZE = {True: (64, 64),
            False: (192, 192)}
    T = 15  # the longest length of text
    vocab_file = 'data/vocab.txt'

    captions_base_dir = os.path.join('data', 'captions')
    categories = os.listdir(captions_base_dir)
    categories.sort()

    if wild_cate not in categories:
        wild_cate = categories[2]

    # Roll out the parameters
    batch_size = 1
    ckpt_dir = Config.ckpt_dir
    results_dir = Config.results_dir
    data_format = Config.data_format
    distance_map = Config.distance_map
    small_img = Config.small_img
    LSTM_hybrid = Config.LSTM_hybrid
    block_type = Config.block_type
    vocab_size = Config.vocab_size

    distance_map = distance_map != 0
    small = small_img != 0
    LSTM_hybrid = LSTM_hybrid != 0

    img_dim = SIZE[small]

    output_folder = results_dir
    print('output_folder:', output_folder)
    os.makedirs(output_folder, exist_ok=True)

    vocab_dict = load_vocab_dict_from_file(vocab_file)

    input_images = tf.placeholder(tf.float32, shape=[1, 3, img_dim[0], img_dim[1]])  # [1, 3, H, W]
    class_ids = tf.placeholder(tf.int32, shape=(1,))  # (1, )
    text_vocab_indiceses = tf.placeholder(tf.int32, shape=[1, 15])  # [1, 15]

    ret_list = build_single_graph(input_images, input_images, None,
                                  class_ids, None,
                                  text_vocab_indiceses,
                                  batch_size=batch_size, training=False,
                                  LSTM_hybrid=LSTM_hybrid,
                                  vocab_size=vocab_size,
                                  data_format=data_format,
                                  distance_map=distance_map,
                                  block_type=block_type)  # [image_gens, images, sketches]

    snapshot_loader = tf.train.Saver()

    tf_config = tf.ConfigProto(allow_soft_placement=True)
    tf_config.gpu_options.allow_growth = True
    with tf.Session(config=tf_config) as sess:
        sess.run(tf.global_variables_initializer())
        sess.run(tf.local_variables_initializer())

        print('Restore trained model:', tf.train.latest_checkpoint(ckpt_dir))
        snapshot_loader.restore(sess, tf.train.latest_checkpoint(ckpt_dir))

        input_name = img_name
        input_category = wild_cate  # e.g. 'bus'
        input_text = wild_text  # e.g. 'A yellow bus with blue window'

        sketch_path = os.path.join(wild_data_base_dir, input_name)
        sketch_image = Image.open(sketch_path)
        sketch_image = sketch_image.convert("RGB")
        # Resize
        if sketch_image.width != img_dim[0] or sketch_image.height != img_dim[1]:
            margin_size = 0 if input_category in ['road'] else 10
            sketch_image = resize_and_padding_mask_image(sketch_image, img_dim[0],
                                                         margin_size=margin_size).astype(np.float32)
        else:
            sketch_image = np.array(sketch_image, dtype=np.float32)  # shape = [H, W, 3]

        # Normalization
        sketch_image = sketch_image / 255.
        sketch_image = sketch_image * 2. - 1

        sketch_image = np.expand_dims(sketch_image, axis=0)  # shape = [1, H, W, 3]
        sketch_image = np.transpose(sketch_image, [0, 3, 1, 2])  # shape = [1, 3, H, W]

        class_id = categories.index(input_category)
        class_id = np.array([class_id])

        vocab_indices = preprocess_sentence(input_text, vocab_dict, T)  # list
        vocab_indices = np.array(vocab_indices, dtype=np.int32)
        vocab_indices = np.expand_dims(vocab_indices, axis=0)  # shape = [1, 15]

        try:
            # print('class_id', class_id)
            # print('vocab_indices', vocab_indices)
            generated_img, _, input_sketch = sess.run(
                [ret_list[0], ret_list[1], ret_list[2]],
                feed_dict={input_images: sketch_image,
                           class_ids: class_id,
                           text_vocab_indiceses: vocab_indices})
        except Exception as e:
            print(e.args)

        if data_format == 'NCHW':
            generated_img = np.transpose(generated_img, (0, 2, 3, 1))
            input_sketch = np.transpose(input_sketch, (0, 2, 3, 1))

        # log('before, generated_img', generated_img)
        # log('before, input_sketch', input_sketch)
        generated_img = ((generated_img + 1) / 2.) * 255
        input_sketch = ((input_sketch + 1) / 2.) * 255
        generated_img = generated_img[:, :, :, ::-1].astype(np.uint8)
        input_sketch = input_sketch.astype(np.uint8)
        # log('after, generated_img', generated_img)
        # log('after, input_sketch', input_sketch)

        img_out_filename = input_name[:-4] + '_output.png'
        sketch_in_filename = input_name[:-4] + '_input.png'

        # Save file
        cv2.imwrite(os.path.join(output_folder, img_out_filename), generated_img[0])
        cv2.imwrite(os.path.join(output_folder, sketch_in_filename), input_sketch[0])

        print('Saved file %s' % img_out_filename)

import os
from time import time
import pickle

import cv2
import numpy as np
from PIL import Image
import scipy.misc
import json
import sys
import tensorflow as tf
from tensorflow.python.client import timeline

from graph_single import build_multi_tower_graph, build_single_graph
from input_pipeline import build_input_queue_paired, build_input_queue_paired_test, resize_and_padding_mask_image, \
    thicken_drawings
from text_processing import preprocess_sentence, load_vocab_dict_from_file
from config import Config

tf.logging.set_verbosity(tf.logging.INFO)
inception_v4_ckpt_path = 'model/inception-cartoon'


def log(name, arr):
    print(name, ', ', arr.shape, ', max:', np.max(arr), ', min:', np.min(arr))


def print_parameter_count(verbose=False):
    total_parameters = 0
    for variable in tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, scope='generator'):
        # shape is an array of tf.Dimension
        shape = variable.get_shape()
        # print(len(shape))
        variable_parametes = 1
        for dim in shape:
            # print(dim)
            variable_parametes *= dim.value
        if verbose and len(shape) > 1:
            print(shape)
            print(variable_parametes)
        total_parameters += variable_parametes
    print('generator')
    print('total_parameters', total_parameters)

    total_parameters = 0
    for variable in tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, scope='discriminator'):
        # shape is an array of tf.Dimension
        shape = variable.get_shape()
        # print(len(shape))
        variable_parametes = 1
        for dim in shape:
            # print(dim)
            variable_parametes *= dim.value
        if verbose and len(shape) > 1:
            print(shape)
            print(variable_parametes)
        total_parameters += variable_parametes
    print('discriminator')
    print('total_parameters', total_parameters)


def train(**kwargs):
    status = 0

    # Roll out the parameters
    appendix = Config.resume_from
    batch_size = Config.batch_size
    max_iter_step = Config.max_iter_step
    Diters = Config.disc_iterations
    ld = Config.ld
    optimizer = Config.optimizer
    lr_G = Config.lr_G
    lr_D = Config.lr_D
    num_gpu = Config.num_gpu
    log_dir = Config.log_dir
    ckpt_dir = Config.ckpt_dir
    data_format = Config.data_format
    distance_map = Config.distance_map
    small_img = Config.small_img
    LSTM_hybrid = Config.LSTM_hybrid
    block_type = Config.block_type
    summary_write_freq = Config.summary_write_freq
    save_model_freq = Config.save_model_freq
    count_left_time_freq = Config.count_left_time_freq
    # count_inception_score_freq = Config.count_inception_score_freq
    vocab_size = Config.vocab_size

    distance_map = distance_map != 0
    small = small_img != 0
    LSTM_hybrid = LSTM_hybrid != 0
    batch_portion = np.array([1, 1, 1, 1], dtype=np.int32)

    iter_from = kwargs['iter_from']

    # Time counter
    prev_time = float("-inf")
    curr_time = float("-inf")

    tf.reset_default_graph()
    print('Iteration starts from: %d' % iter_from)

    # assert inception_score.softmax.graph != tf.get_default_graph()
    # inception_score._init_inception()

    counter = tf.Variable(initial_value=iter_from, dtype=tf.int32, trainable=False)
    counter_addition_op = tf.assign_add(counter, 1, use_locking=True)

    # Construct data queue
    with tf.device('/cpu:0'):
        images, sketches, image_paired_class_ids, _, _, _, text_vocab_indiceses = build_input_queue_paired(
            mode='train',
            batch_size=batch_size * num_gpu,
            data_format=data_format,
            distance_map=distance_map,
            small=small, capacity=2 ** 11)  # images/sketches [2, 3, H, W], text_vocab_indiceses [2, 15]
    with tf.device('/cpu:0'):
        images_d, _, image_paired_class_ids_d, _, _, _, _ = build_input_queue_paired(
            mode='train',
            batch_size=batch_size * num_gpu,
            data_format=data_format,
            distance_map=distance_map,
            small=small, capacity=2 ** 11)  # [2, 3, H, W]

    opt_g, opt_d, loss_g, loss_d, merged_all = build_multi_tower_graph(
        images, sketches, images_d,
        image_paired_class_ids, image_paired_class_ids_d,
        text_vocab_indiceses,
        LSTM_hybrid=LSTM_hybrid,
        vocab_size=vocab_size,
        batch_size=batch_size, num_gpu=num_gpu, batch_portion=batch_portion, training=True,
        learning_rates={
            "generator": lr_G,
            "discriminator": lr_D,
        },
        counter=counter, max_iter_step=max_iter_step,
        ld=ld, data_format=data_format,
        distance_map=distance_map,
        optimizer=optimizer,
        block_type=block_type)

    saver = tf.train.Saver(max_to_keep=100)
    # try:
    #     inception_loader = tf.train.Saver(tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope='InceptionV4'))
    #     perceptual_model_checkpoint_path = inception_v4_ckpt_path
    #     perceptual_model_path = tf.train.latest_checkpoint(perceptual_model_checkpoint_path)
    # except:
    #     inception_loader = None

    config = tf.ConfigProto(allow_soft_placement=True,
                            intra_op_parallelism_threads=4, inter_op_parallelism_threads=4)
    config.gpu_options.allow_growth = True

    with tf.Session(config=config) as sess:
        sess.run(tf.global_variables_initializer())
        sess.run(tf.local_variables_initializer())

        # if inception_loader is not None:
        #     print('Restore:', perceptual_model_path)
        #     inception_loader.restore(sess, perceptual_model_path)

        summary_writer = tf.summary.FileWriter(log_dir, sess.graph)
        if iter_from > 0:
            snapshot_loader = tf.train.Saver()
            print('Restore:', tf.train.latest_checkpoint(ckpt_dir))
            snapshot_loader.restore(sess, tf.train.latest_checkpoint(ckpt_dir))
            summary_writer.reopen()

        run_options = tf.RunOptions(trace_level=tf.RunOptions.NO_TRACE)
        run_metadata = tf.RunMetadata()

        print_parameter_count(verbose=False)

        coord = tf.train.Coordinator()
        threads = tf.train.start_queue_runners(sess=sess, coord=coord)

        sess.run([counter.assign(iter_from)])

        for i in range(iter_from, max_iter_step):
            if status == -1:
                break

            ## count left time
            if i % count_left_time_freq == 0:
                curr_time = time()
                elapsed = curr_time - prev_time
                print("Now at iteration %d. Elapsed time: %.5fs. Average time: %.5fs/iter"
                      % (i, elapsed, elapsed / 100.))

                if elapsed != float("inf"):
                    left_iter = max_iter_step - i
                    left_sec = left_iter * (elapsed / 100.)
                    left_day = int(left_sec / 24 / 60 / 60)
                    left_hour = int((left_sec - (24 * 60 * 60) * left_day) / 60 / 60)
                    left_min = int((left_sec - (24 * 60 * 60) * left_day - (60 * 60) * left_hour) / 60)
                    print("Left time:%dd %dh %dm" % (left_day, left_hour, left_min))

                prev_time = curr_time

            diters = Diters

            # Train Discriminator
            for j in range(diters):
                ## summary
                if i % summary_write_freq == 0 and j == 0:
                    _, merged, loss_d_out = sess.run([opt_d, merged_all, loss_d],
                                                     options=run_options,
                                                     run_metadata=run_metadata)
                    summary_writer.add_summary(merged, i)
                    summary_writer.flush()
                else:
                    _, loss_d_out = sess.run([opt_d, loss_d])
                # print('loss_d', loss_d_out)
                if np.isnan(np.sum(loss_d_out)):
                    status = -1
                    print("NaN occurred during training D")
                    return status

            # Train Generator
            if i % summary_write_freq == 0:
                _, merged, loss_g_out, counter_out, _ = sess.run(
                    [opt_g, merged_all, loss_g, counter, counter_addition_op],
                    options=run_options,
                    run_metadata=run_metadata)
                summary_writer.add_summary(merged, i)
                summary_writer.flush()
            else:
                _, loss_g_out, counter_out, _ = sess.run([opt_g, loss_g, counter, counter_addition_op])
            # print('loss_g', loss_g_out)
            if np.isnan(np.sum(loss_g_out)):
                status = -1
                print("NaN occurred during training G")
                return status

            ## save model
            if i % save_model_freq == save_model_freq - 1:
                saver.save(sess, os.path.join(ckpt_dir, 'model_{}.ckpt'.format(i)), global_step=i)
                print('Save model_{}.ckpt'.format(i))

        coord.request_stop()
        coord.join(threads)

    return status


def validation(**kwargs):
    # Roll out the parameters
    dataset_type = Config.dataset_type
    batch_size = Config.batch_size
    ckpt_dir = Config.ckpt_dir
    results_dir = Config.results_dir
    data_format = Config.data_format
    distance_map = Config.distance_map
    small_img = Config.small_img
    LSTM_hybrid = Config.LSTM_hybrid
    block_type = Config.block_type
    vocab_size = Config.vocab_size

    channel = 3
    distance_map = distance_map != 0
    small = small_img != 0
    LSTM_hybrid = LSTM_hybrid != 0

    if LSTM_hybrid:
        output_folder = os.path.join(results_dir, 'with_text')
    else:
        output_folder = os.path.join(results_dir, 'without_text')
    print(output_folder)
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)

    # Time counter
    prev_time = float("-inf")
    curr_time = float("-inf")
    # Construct data queue
    with tf.device('/cpu:0'):
        images, sketches, class_ids, categories, image_names, color_texts, text_vocab_indiceses \
            = build_input_queue_paired_test(
            mode=dataset_type,
            batch_size=batch_size, data_format=data_format,
            distance_map=distance_map, small=small, capacity=512)  # [2, 3, H, W]

    ret_list = build_single_graph(images, sketches, None,
                                  class_ids, None,
                                  text_vocab_indiceses,
                                  batch_size=batch_size, training=False,
                                  LSTM_hybrid=LSTM_hybrid,
                                  vocab_size=vocab_size,
                                  data_format=data_format,
                                  distance_map=distance_map,
                                  block_type=block_type)  # [image_gens, images, sketches]

    snapshot_loader = tf.train.Saver()

    tf_config = tf.ConfigProto(allow_soft_placement=True)
    tf_config.gpu_options.allow_growth = True
    with tf.Session(config=tf_config) as sess:
        sess.run(tf.global_variables_initializer())
        sess.run(tf.local_variables_initializer())

        print('Restore trained model:', tf.train.latest_checkpoint(ckpt_dir))
        snapshot_loader.restore(sess, tf.train.latest_checkpoint(ckpt_dir))
        counter = 0

        coord = tf.train.Coordinator()
        threads = tf.train.start_queue_runners(sess=sess, coord=coord)

        while True:
            try:
                generated_img, gt_image, input_sketch, category, image_name = sess.run(
                    [ret_list[0], ret_list[1], ret_list[2], categories, image_names])
            except Exception as e:
                print(e.args)
                break

            if counter % 100 == 0:
                curr_time = time()
                elapsed = curr_time - prev_time
                print(
                    "Now at iteration %d. Elapsed time: %.5fs." % (counter, elapsed))
                prev_time = curr_time

            if data_format == 'NCHW':
                generated_img = np.transpose(generated_img, (0, 2, 3, 1))
                gt_image = np.transpose(gt_image, (0, 2, 3, 1))
                input_sketch = np.transpose(input_sketch, (0, 2, 3, 1))

            # log('before, generated_img', generated_img)
            # log('before, gt_image', gt_image)
            # log('before, input_sketch', input_sketch)
            generated_img = ((generated_img + 1) / 2.) * 255
            gt_image = ((gt_image + 1) / 2.) * 255
            input_sketch = ((input_sketch + 1) / 2.) * 255
            generated_img = generated_img[:, :, :, ::-1].astype(np.uint8)
            gt_image = gt_image[:, :, :, ::-1].astype(np.uint8)
            input_sketch = input_sketch.astype(np.uint8)
            # log('after, generated_img', generated_img)
            # log('after, gt_image', gt_image)  # (2, H, W, 3)
            # log('after, input_sketch', input_sketch)

            for i in range(batch_size):
                this_prefix = '%s' % (category[i].decode('ascii'))
                img_out_filename = this_prefix + '_' + image_name[i].decode()[:-4] + '_output.png'
                img_gt_filename = this_prefix + '_' + image_name[i].decode()[:-4] + '_target.png'
                sketch_in_filename = this_prefix + '_' + image_name[i].decode()[:-4] + '_input.png'

                # Save file
                # file_path = os.path.join(output_folder, 'output_%d.jpg' % int(counter / batch_size))
                cv2.imwrite(os.path.join(output_folder, img_out_filename), generated_img[i])
                cv2.imwrite(os.path.join(output_folder, img_gt_filename), gt_image[i])
                cv2.imwrite(os.path.join(output_folder, sketch_in_filename), input_sketch[i])
                # output_img = np.zeros((img_dim * 2, img_dim * batch_size, channel))

                print('Saved file %s' % this_prefix)

            counter += 1

        coord.request_stop()
        coord.join(threads)


def test():
    SIZE = {True: (64, 64),
            False: (192, 192)}
    T = 15  # the longest length of text
    vocab_file = 'data/vocab.txt'
    test_data_base_dir = 'data'

    captions_base_dir = os.path.join(test_data_base_dir, 'captions')
    images_base_dir = os.path.join(test_data_base_dir, 'images')
    categories = os.listdir(captions_base_dir)
    categories.sort()
    print(categories)

    # Roll out the parameters
    batch_size = 1
    ckpt_dir = Config.ckpt_dir
    results_dir = Config.results_dir
    data_format = Config.data_format
    distance_map = Config.distance_map
    small_img = Config.small_img
    LSTM_hybrid = Config.LSTM_hybrid
    block_type = Config.block_type
    vocab_size = Config.vocab_size

    distance_map = distance_map != 0
    small = small_img != 0
    LSTM_hybrid = LSTM_hybrid != 0

    img_dim = SIZE[small]

    output_folder = results_dir
    os.makedirs(output_folder, exist_ok=True)

    vocab_dict = load_vocab_dict_from_file(vocab_file)

    input_images = tf.placeholder(tf.float32, shape=[1, 3, img_dim[0], img_dim[1]])  # [1, 3, H, W]
    class_ids = tf.placeholder(tf.int32, shape=(1,))  # (1, )
    text_vocab_indiceses = tf.placeholder(tf.int32, shape=[1, 15])  # [1, 15]

    ret_list = build_single_graph(input_images, input_images, None,
                                  class_ids, None,
                                  text_vocab_indiceses,
                                  batch_size=batch_size, training=False,
                                  LSTM_hybrid=LSTM_hybrid,
                                  vocab_size=vocab_size,
                                  data_format=data_format,
                                  distance_map=distance_map,
                                  block_type=block_type)  # [image_gens, images, sketches]

    snapshot_loader = tf.train.Saver()

    tf_config = tf.ConfigProto(allow_soft_placement=True)
    tf_config.gpu_options.allow_growth = True
    with tf.Session(config=tf_config) as sess:
        sess.run(tf.global_variables_initializer())
        sess.run(tf.local_variables_initializer())

        print('Restore trained model:', tf.train.latest_checkpoint(ckpt_dir))
        snapshot_loader.restore(sess, tf.train.latest_checkpoint(ckpt_dir))

        for cate in categories:
            testing_json = os.path.join(captions_base_dir, cate, 'test.json')
            fp = open(testing_json, "r")
            json_data = fp.read()
            json_data = json.loads(json_data)
            print(len(json_data), 'inference datas')

            for i in range(len(json_data)):
                input_name = json_data[i]['key']  # e.g. '228_1.png'
                input_text = json_data[i]['color_text']  # e.g. 'A yellow bus with blue window'

                sketch_path = os.path.join(images_base_dir, cate, 'sketch', input_name)
                sketch_image = Image.open(sketch_path)
                sketch_image = sketch_image.convert("RGB")
                # Resize
                if sketch_image.width != img_dim[0] or sketch_image.height != img_dim[1]:
                    margin_size = 0 if cate in ['road'] else 10
                    sketch_image = resize_and_padding_mask_image(sketch_image, img_dim[0],
                                                                 margin_size=margin_size).astype(np.float32)
                else:
                    sketch_image = np.array(sketch_image, dtype=np.float32)  # shape = [H, W, 3]

                if cate in ['house', 'road']:
                    sketch_image = thicken_drawings(sketch_image).astype(np.float32)  # shape = [H, W, 3]

                # Normalization
                sketch_image = sketch_image / 255.
                sketch_image = sketch_image * 2. - 1

                sketch_image = np.expand_dims(sketch_image, axis=0)  # shape = [1, H, W, 3]
                sketch_image = np.transpose(sketch_image, [0, 3, 1, 2])  # shape = [1, 3, H, W]

                class_id = categories.index(cate)
                class_id = np.array([class_id])

                vocab_indices = preprocess_sentence(input_text, vocab_dict, T)  # list
                vocab_indices = np.array(vocab_indices, dtype=np.int32)
                vocab_indices = np.expand_dims(vocab_indices, axis=0)  # shape = [1, 15]

                try:
                    # print('class_id', class_id)
                    # print('vocab_indices', vocab_indices)
                    generated_img, _, input_sketch = sess.run(
                        [ret_list[0], ret_list[1], ret_list[2]],
                        feed_dict={input_images: sketch_image,
                                   class_ids: class_id,
                                   text_vocab_indiceses: vocab_indices})
                except Exception as e:
                    print(e.args)
                    break

                if data_format == 'NCHW':
                    generated_img = np.transpose(generated_img, (0, 2, 3, 1))
                    input_sketch = np.transpose(input_sketch, (0, 2, 3, 1))

                # log('before, generated_img', generated_img)
                # log('before, input_sketch', input_sketch)
                generated_img = ((generated_img + 1) / 2.) * 255
                input_sketch = ((input_sketch + 1) / 2.) * 255
                generated_img = generated_img[:, :, :, ::-1].astype(np.uint8)
                input_sketch = input_sketch.astype(np.uint8)
                # log('after, generated_img', generated_img)
                # log('after, input_sketch', input_sketch)

                img_out_filename = cate + '_' + input_name[:-4] + '_output.png'
                sketch_in_filename = cate + '_' + input_name[:-4] + '_input.png'

                # Save file
                cv2.imwrite(os.path.join(output_folder, img_out_filename), generated_img[0])
                cv2.imwrite(os.path.join(output_folder, sketch_in_filename), input_sketch[0])

                print('Saved file %s' % img_out_filename)


def inference(img_name, instruction):
    wild_data_base_dir = 'examples'
    wild_text = instruction

    wild_cate = img_name[:img_name.find('.png')]

    SIZE = {True: (64, 64),
            False: (192, 192)}
    T = 15  # the longest length of text
    vocab_file = 'data/vocab.txt'

    captions_base_dir = os.path.join('data', 'captions')
    categories = os.listdir(captions_base_dir)
    categories.sort()

    if wild_cate not in categories:
        wild_cate = categories[2]

    # Roll out the parameters
    batch_size = 1
    ckpt_dir = Config.ckpt_dir
    results_dir = Config.results_dir
    data_format = Config.data_format
    distance_map = Config.distance_map
    small_img = Config.small_img
    LSTM_hybrid = Config.LSTM_hybrid
    block_type = Config.block_type
    vocab_size = Config.vocab_size

    distance_map = distance_map != 0
    small = small_img != 0
    LSTM_hybrid = LSTM_hybrid != 0

    img_dim = SIZE[small]

    output_folder = results_dir
    print('output_folder:', output_folder)
    os.makedirs(output_folder, exist_ok=True)

    vocab_dict = load_vocab_dict_from_file(vocab_file)

    input_images = tf.placeholder(tf.float32, shape=[1, 3, img_dim[0], img_dim[1]])  # [1, 3, H, W]
    class_ids = tf.placeholder(tf.int32, shape=(1,))  # (1, )
    text_vocab_indiceses = tf.placeholder(tf.int32, shape=[1, 15])  # [1, 15]

    ret_list = build_single_graph(input_images, input_images, None,
                                  class_ids, None,
                                  text_vocab_indiceses,
                                  batch_size=batch_size, training=False,
                                  LSTM_hybrid=LSTM_hybrid,
                                  vocab_size=vocab_size,
                                  data_format=data_format,
                                  distance_map=distance_map,
                                  block_type=block_type)  # [image_gens, images, sketches]

    snapshot_loader = tf.train.Saver()

    tf_config = tf.ConfigProto(allow_soft_placement=True)
    tf_config.gpu_options.allow_growth = True
    with tf.Session(config=tf_config) as sess:
        sess.run(tf.global_variables_initializer())
        sess.run(tf.local_variables_initializer())

        print('Restore trained model:', tf.train.latest_checkpoint(ckpt_dir))
        snapshot_loader.restore(sess, tf.train.latest_checkpoint(ckpt_dir))

        input_name = img_name
        input_category = wild_cate  # e.g. 'bus'
        input_text = wild_text  # e.g. 'A yellow bus with blue window'

        sketch_path = os.path.join(wild_data_base_dir, input_name)
        sketch_image = Image.open(sketch_path)
        sketch_image = sketch_image.convert("RGB")
        # Resize
        if sketch_image.width != img_dim[0] or sketch_image.height != img_dim[1]:
            margin_size = 0 if input_category in ['road'] else 10
            sketch_image = resize_and_padding_mask_image(sketch_image, img_dim[0],
                                                         margin_size=margin_size).astype(np.float32)
        else:
            sketch_image = np.array(sketch_image, dtype=np.float32)  # shape = [H, W, 3]

        # Normalization
        sketch_image = sketch_image / 255.
        sketch_image = sketch_image * 2. - 1

        sketch_image = np.expand_dims(sketch_image, axis=0)  # shape = [1, H, W, 3]
        sketch_image = np.transpose(sketch_image, [0, 3, 1, 2])  # shape = [1, 3, H, W]

        class_id = categories.index(input_category)
        class_id = np.array([class_id])

        vocab_indices = preprocess_sentence(input_text, vocab_dict, T)  # list
        vocab_indices = np.array(vocab_indices, dtype=np.int32)
        vocab_indices = np.expand_dims(vocab_indices, axis=0)  # shape = [1, 15]

        try:
            # print('class_id', class_id)
            # print('vocab_indices', vocab_indices)
            generated_img, _, input_sketch = sess.run(
                [ret_list[0], ret_list[1], ret_list[2]],
                feed_dict={input_images: sketch_image,
                           class_ids: class_id,
                           text_vocab_indiceses: vocab_indices})
        except Exception as e:
            print(e.args)

        if data_format == 'NCHW':
            generated_img = np.transpose(generated_img, (0, 2, 3, 1))
            input_sketch = np.transpose(input_sketch, (0, 2, 3, 1))

        # log('before, generated_img', generated_img)
        # log('before, input_sketch', input_sketch)
        generated_img = ((generated_img + 1) / 2.) * 255
        input_sketch = ((input_sketch + 1) / 2.) * 255
        generated_img = generated_img[:, :, :, ::-1].astype(np.uint8)
        input_sketch = input_sketch.astype(np.uint8)
        # log('after, generated_img', generated_img)
        # log('after, input_sketch', input_sketch)

        img_out_filename = input_name[:-4] + '_output.png'
        sketch_in_filename = input_name[:-4] + '_input.png'

        # Save file
        cv2.imwrite(os.path.join(output_folder, img_out_filename), generated_img[0])
        cv2.imwrite(os.path.join(output_folder, sketch_in_filename), input_sketch[0])

        print('Saved file %s' % img_out_filename)

# shape is an array of tf.Dimension # print(len(shape)) # print(dim) # shape is an array of tf.Dimension # print(len(shape)) # print(dim) # Roll out the parameters # count_inception_score_freq = Config.count_inception_score_freq # Time counter # assert inception_score.softmax.graph != tf.get_default_graph() # inception_score._init_inception() # Construct data queue # images/sketches [2, 3, H, W], text_vocab_indiceses [2, 15] # [2, 3, H, W] # try: #     inception_loader = tf.train.Saver(tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope='InceptionV4')) #     perceptual_model_checkpoint_path = inception_v4_ckpt_path #     perceptual_model_path = tf.train.latest_checkpoint(perceptual_model_checkpoint_path) # except: #     inception_loader = None # if inception_loader is not None: #     print('Restore:', perceptual_model_path) #     inception_loader.restore(sess, perceptual_model_path) ## count left time # Train Discriminator ## summary # print('loss_d', loss_d_out) # Train Generator # print('loss_g', loss_g_out) ## save model # Roll out the parameters # Time counter # Construct data queue # [2, 3, H, W] # [image_gens, images, sketches] # log('before, generated_img', generated_img) # log('before, gt_image', gt_image) # log('before, input_sketch', input_sketch) # log('after, generated_img', generated_img) # log('after, gt_image', gt_image)  # (2, H, W, 3) # log('after, input_sketch', input_sketch) # Save file # file_path = os.path.join(output_folder, 'output_%d.jpg' % int(counter / batch_size)) # output_img = np.zeros((img_dim * 2, img_dim * batch_size, channel)) # the longest length of text # Roll out the parameters # [1, 3, H, W] # (1, ) # [1, 15] # [image_gens, images, sketches] # e.g. '228_1.png' # e.g. 'A yellow bus with blue window' # Resize # shape = [H, W, 3] # shape = [H, W, 3] # Normalization # shape = [1, H, W, 3] # shape = [1, 3, H, W] # list # shape = [1, 15] # print('class_id', class_id) # print('vocab_indices', vocab_indices) # log('before, generated_img', generated_img) # log('before, input_sketch', input_sketch) # log('after, generated_img', generated_img) # log('after, input_sketch', input_sketch) # Save file # the longest length of text # Roll out the parameters # [1, 3, H, W] # (1, ) # [1, 15] # [image_gens, images, sketches] # e.g. 'bus' # e.g. 'A yellow bus with blue window' # Resize # shape = [H, W, 3] # Normalization # shape = [1, H, W, 3] # shape = [1, 3, H, W] # list # shape = [1, 15] # print('class_id', class_id) # print('vocab_indices', vocab_indices) # log('before, generated_img', generated_img) # log('before, input_sketch', input_sketch) # log('after, generated_img', generated_img) # log('after, input_sketch', input_sketch) # Save file