<commit_before>""" Contains the definition of the ChangeGAN architecture. """
import multiprocessing

import tensorflow as tf

from gan_utils import encoder, decoder, transformer, discriminator, preprocess_image

slim = tf.contrib.slim

default_image_size = 256


def model_fn(inputs_a, inputs_b, learning_rate, num_blocks=9, is_training=True, scope=None, weight_decay=0.0001):
    encoder_dims = [32, 64, 128]
    deep_encoder_dims = [64, 128, 256]
    decoder_dims = [64, 32, 3]
    deep_decoder_dims = [128, 64, 3]
    with tf.variable_scope(scope, 'ChangeGAN', [inputs_a, inputs_b]):
        with slim.arg_scope([slim.conv2d, slim.conv2d_transpose],
                            weights_regularizer=slim.l2_regularizer(weight_decay)):
            with slim.arg_scope([slim.batch_norm],
                                is_training=is_training):
                def converter_ab(inputs_a, reuse=None):
                    ################
                    # Encoder part #
                    ################
                    z_a = encoder(inputs_a, deep_encoder_dims, scope='Encoder_A', reuse=reuse)

                    # z_a is split into c_b, z_a-b
                    c_b, z_a_b = tf.split(z_a, num_or_size_splits=2, axis=3)

                    ####################
                    # Transformer part #
                    ####################
                    c_b = transformer(c_b, encoder_dims[-1], num_blocks=num_blocks,
                                      scope='Transformer_B', reuse=reuse)

                    ################
                    # Decoder part #
                    ################
                    outputs_b = decoder(c_b, decoder_dims, scope='Decoder_B', reuse=reuse)

                    return outputs_b, z_a_b

                def converter_ba(inputs_b, z_a_b, reuse=None):
                    ################
                    # Encoder part #
                    ################
                    z_b = encoder(inputs_b, encoder_dims, scope='Encoder_B', reuse=reuse)

                    # Concat z_b and z_a-b
                    c_a = tf.concat([z_b, z_a_b], 3)

                    ####################
                    # Transformer part #
                    ####################
                    c_a = transformer(c_a, deep_encoder_dims[-1], num_blocks=num_blocks,
                                      scope='Transformer_A', reuse=reuse)

                    ################
                    # Decoder part #
                    ################
                    outputs_a = decoder(c_a, deep_decoder_dims, scope='Decoder_A', reuse=reuse)

                    return outputs_a

                bbox_channel_a = _get_bbox(inputs_a)

                outputs_ab, z_a_b = converter_ab(inputs_a)
                outputs_bbox_ab = tf.concat([outputs_ab, bbox_channel_a], 3)
                outputs_ba = converter_ba(inputs_b, z_a_b)
                outputs_bbox_ba = tf.concat([outputs_ba, bbox_channel_a], 3)

                outputs_bab, _ = converter_ab(outputs_bbox_ba, reuse=True)
                outputs_aba = converter_ba(outputs_bbox_ab, z_a_b, reuse=True)

                logits_a_real, probs_a_real = discriminator(inputs_a, deep_encoder_dims, scope='Discriminator_A')
                logits_a_fake, probs_a_fake = discriminator(outputs_bbox_ba, deep_encoder_dims, scope='Discriminator_A', reuse=True)
                logits_b_real, probs_b_real = discriminator(inputs_b, deep_encoder_dims, scope='Discriminator_B')
                logits_b_fake, probs_b_fake = discriminator(outputs_bbox_ab, deep_encoder_dims, scope='Discriminator_B', reuse=True)

                outputs = [outputs_ba, outputs_ab, outputs_aba, outputs_bab]

    with tf.name_scope('images'):
        tf.summary.image('X_A', _remove_bbox(inputs_a))
        tf.summary.image('X_B', _remove_bbox(inputs_b))
        tf.summary.image('X_BA', outputs_ba)
        tf.summary.image('X_AB', outputs_ab)
        tf.summary.image('X_ABA', outputs_aba)
        tf.summary.image('X_BAB', outputs_bab)

    global_step = tf.train.get_or_create_global_step()

    if not is_training:
        return outputs

    t_vars = tf.trainable_variables()

    d_a_vars = [var for var in t_vars if 'Discriminator_A' in var.name]
    d_b_vars = [var for var in t_vars if 'Discriminator_B' in var.name]
    g_vars = [var for var in t_vars if 'coder' in var.name or 'Transformer' in var.name]

    ##########
    # Losses #
    ##########
    # Losses for discriminator
    l_d_a_fake = tf.reduce_mean(tf.square(logits_a_fake))
    l_d_a_real = tf.reduce_mean(tf.squared_difference(logits_a_real, 1.))
    l_d_a = 0.5 * (l_d_a_fake + l_d_a_real)
    train_op_d_a = tf.train.AdamOptimizer(
        learning_rate=learning_rate,
        beta1=0.5,
        beta2=0.999
    ).minimize(l_d_a, global_step=global_step, var_list=d_a_vars)

    l_d_b_fake = tf.reduce_mean(tf.square(logits_b_fake))
    l_d_b_real = tf.reduce_mean(tf.squared_difference(logits_b_real, 1.))
    l_d_b = 0.5 * (l_d_b_fake + l_d_b_real)
    train_op_d_b = tf.train.AdamOptimizer(
        learning_rate=learning_rate,
        beta1=0.5,
        beta2=0.999
    ).minimize(l_d_b, global_step=global_step, var_list=d_b_vars)

    l_d = l_d_a + l_d_b

    # Losses for generators
    l_g_a = tf.reduce_mean(tf.squared_difference(logits_a_fake, 1.))
    l_g_b = tf.reduce_mean(tf.squared_difference(logits_b_fake, 1.))
    l_const_a = tf.reduce_mean(tf.losses.absolute_difference(_remove_bbox(inputs_a), outputs_aba))
    l_const_b = tf.reduce_mean(tf.losses.absolute_difference(_remove_bbox(inputs_b), outputs_bab))

    l_g = l_g_a + l_g_b + 10. * (l_const_a + l_const_b)
    train_op_g = tf.train.AdamOptimizer(
        learning_rate=learning_rate,
        beta1=0.5,
        beta2=0.999
    ).minimize(l_g, global_step=global_step, var_list=g_vars)

    with tf.name_scope('losses'):
        tf.summary.scalar('L_D_A_Fake', l_d_a_fake)
        tf.summary.scalar('L_D_A_Real', l_d_a_real)
        tf.summary.scalar('L_D_A', l_d_a)
        tf.summary.scalar('L_D_B_Fake', l_d_b_fake)
        tf.summary.scalar('L_D_B_Real', l_d_b_real)
        tf.summary.scalar('L_D_B', l_d_b)
        tf.summary.scalar('L_D', l_d)

        tf.summary.scalar('L_G_A', l_g_a)
        tf.summary.scalar('L_G_B', l_g_b)
        tf.summary.scalar('L_Const_A', l_const_a)
        tf.summary.scalar('L_Const_B', l_const_b)
        tf.summary.scalar('L_G', l_g)

    train_op = tf.group(*[train_op_d_a, train_op_d_b, train_op_g])

    return train_op, global_step, outputs


def input_fn(dataset_a, dataset_b, batch_size=1, num_readers=4, is_training=True):
    provider_a = slim.dataset_data_provider.DatasetDataProvider(
        dataset_a,
        num_readers=num_readers,
        common_queue_capacity=20 * batch_size,
        common_queue_min=10 * batch_size)
    provider_b = slim.dataset_data_provider.DatasetDataProvider(
        dataset_b,
        num_readers=num_readers,
        common_queue_capacity=20 * batch_size,
        common_queue_min=10 * batch_size)
    [image_a, bbox_a] = provider_a.get(['image', 'object/bbox'])
    [image_b, bbox_b] = provider_b.get(['image', 'object/bbox'])

    train_image_size = default_image_size

    def add_channel(image, bbox):
        ymin = bbox[0]
        xmin = bbox[1]
        ymax = bbox[2]
        xmax = bbox[3]

        image_shape = tf.to_float(tf.shape(image))
        height = image_shape[0]
        width = image_shape[1]

        bbox_height = (ymax - ymin) * height
        bbox_width = (xmax - xmin) * width
        channel = tf.ones(tf.to_int32(tf.stack([bbox_height, bbox_width])))
        channel = tf.expand_dims(channel, axis=2)

        pad_top = tf.to_int32(ymin * height)
        pad_left = tf.to_int32(xmin * width)
        height = tf.to_int32(height)
        width = tf.to_int32(width)
        channel = tf.image.pad_to_bounding_box(channel, pad_top, pad_left, height, width)
        # TODO: Decide pad one or zero
        channel = tf.ones_like(channel) - channel

        image = tf.concat([image, channel], axis=2)

        return image

    image_a = tf.image.convert_image_dtype(image_a, dtype=tf.float32)
    image_b = tf.image.convert_image_dtype(image_b, dtype=tf.float32)
    # [Num of boxes, 4] => [4]
    bbox_a = tf.squeeze(bbox_a, axis=0)<commit_msg>Zero padding for bound box<commit_after>""" Contains the definition of the ChangeGAN architecture. """
import multiprocessing

import tensorflow as tf

from gan_utils import encoder, decoder, transformer, discriminator, preprocess_image

slim = tf.contrib.slim

default_image_size = 256


def model_fn(inputs_a, inputs_b, learning_rate, num_blocks=9, is_training=True, scope=None, weight_decay=0.0001):
    encoder_dims = [32, 64, 128]
    deep_encoder_dims = [64, 128, 256]
    decoder_dims = [64, 32, 3]
    deep_decoder_dims = [128, 64, 3]
    with tf.variable_scope(scope, 'ChangeGAN', [inputs_a, inputs_b]):
        with slim.arg_scope([slim.conv2d, slim.conv2d_transpose],
                            weights_regularizer=slim.l2_regularizer(weight_decay)):
            with slim.arg_scope([slim.batch_norm],
                                is_training=is_training):
                def converter_ab(inputs_a, reuse=None):
                    ################
                    # Encoder part #
                    ################
                    z_a = encoder(inputs_a, deep_encoder_dims, scope='Encoder_A', reuse=reuse)

                    # z_a is split into c_b, z_a-b
                    c_b, z_a_b = tf.split(z_a, num_or_size_splits=2, axis=3)

                    ####################
                    # Transformer part #
                    ####################
                    c_b = transformer(c_b, encoder_dims[-1], num_blocks=num_blocks,
                                      scope='Transformer_B', reuse=reuse)

                    ################
                    # Decoder part #
                    ################
                    outputs_b = decoder(c_b, decoder_dims, scope='Decoder_B', reuse=reuse)

                    return outputs_b, z_a_b

                def converter_ba(inputs_b, z_a_b, reuse=None):
                    ################
                    # Encoder part #
                    ################
                    z_b = encoder(inputs_b, encoder_dims, scope='Encoder_B', reuse=reuse)

                    # Concat z_b and z_a-b
                    c_a = tf.concat([z_b, z_a_b], 3)

                    ####################
                    # Transformer part #
                    ####################
                    c_a = transformer(c_a, deep_encoder_dims[-1], num_blocks=num_blocks,
                                      scope='Transformer_A', reuse=reuse)

                    ################
                    # Decoder part #
                    ################
                    outputs_a = decoder(c_a, deep_decoder_dims, scope='Decoder_A', reuse=reuse)

                    return outputs_a

                bbox_channel_a = _get_bbox(inputs_a)

                outputs_ab, z_a_b = converter_ab(inputs_a)
                outputs_bbox_ab = tf.concat([outputs_ab, bbox_channel_a], 3)
                outputs_ba = converter_ba(inputs_b, z_a_b)
                outputs_bbox_ba = tf.concat([outputs_ba, bbox_channel_a], 3)

                outputs_bab, _ = converter_ab(outputs_bbox_ba, reuse=True)
                outputs_aba = converter_ba(outputs_bbox_ab, z_a_b, reuse=True)

                logits_a_real, probs_a_real = discriminator(inputs_a, deep_encoder_dims, scope='Discriminator_A')
                logits_a_fake, probs_a_fake = discriminator(outputs_bbox_ba, deep_encoder_dims, scope='Discriminator_A', reuse=True)
                logits_b_real, probs_b_real = discriminator(inputs_b, deep_encoder_dims, scope='Discriminator_B')
                logits_b_fake, probs_b_fake = discriminator(outputs_bbox_ab, deep_encoder_dims, scope='Discriminator_B', reuse=True)

                outputs = [outputs_ba, outputs_ab, outputs_aba, outputs_bab]

    with tf.name_scope('images'):
        tf.summary.image('X_A', _remove_bbox(inputs_a))
        tf.summary.image('X_B', _remove_bbox(inputs_b))
        tf.summary.image('X_BA', outputs_ba)
        tf.summary.image('X_AB', outputs_ab)
        tf.summary.image('X_ABA', outputs_aba)
        tf.summary.image('X_BAB', outputs_bab)

    global_step = tf.train.get_or_create_global_step()

    if not is_training:
        return outputs

    t_vars = tf.trainable_variables()

    d_a_vars = [var for var in t_vars if 'Discriminator_A' in var.name]
    d_b_vars = [var for var in t_vars if 'Discriminator_B' in var.name]
    g_vars = [var for var in t_vars if 'coder' in var.name or 'Transformer' in var.name]

    ##########
    # Losses #
    ##########
    # Losses for discriminator
    l_d_a_fake = tf.reduce_mean(tf.square(logits_a_fake))
    l_d_a_real = tf.reduce_mean(tf.squared_difference(logits_a_real, 1.))
    l_d_a = 0.5 * (l_d_a_fake + l_d_a_real)
    train_op_d_a = tf.train.AdamOptimizer(
        learning_rate=learning_rate,
        beta1=0.5,
        beta2=0.999
    ).minimize(l_d_a, global_step=global_step, var_list=d_a_vars)

    l_d_b_fake = tf.reduce_mean(tf.square(logits_b_fake))
    l_d_b_real = tf.reduce_mean(tf.squared_difference(logits_b_real, 1.))
    l_d_b = 0.5 * (l_d_b_fake + l_d_b_real)
    train_op_d_b = tf.train.AdamOptimizer(
        learning_rate=learning_rate,
        beta1=0.5,
        beta2=0.999
    ).minimize(l_d_b, global_step=global_step, var_list=d_b_vars)

    l_d = l_d_a + l_d_b

    # Losses for generators
    l_g_a = tf.reduce_mean(tf.squared_difference(logits_a_fake, 1.))
    l_g_b = tf.reduce_mean(tf.squared_difference(logits_b_fake, 1.))
    l_const_a = tf.reduce_mean(tf.losses.absolute_difference(_remove_bbox(inputs_a), outputs_aba))
    l_const_b = tf.reduce_mean(tf.losses.absolute_difference(_remove_bbox(inputs_b), outputs_bab))

    l_g = l_g_a + l_g_b + 10. * (l_const_a + l_const_b)
    train_op_g = tf.train.AdamOptimizer(
        learning_rate=learning_rate,
        beta1=0.5,
        beta2=0.999
    ).minimize(l_g, global_step=global_step, var_list=g_vars)

    with tf.name_scope('losses'):
        tf.summary.scalar('L_D_A_Fake', l_d_a_fake)
        tf.summary.scalar('L_D_A_Real', l_d_a_real)
        tf.summary.scalar('L_D_A', l_d_a)
        tf.summary.scalar('L_D_B_Fake', l_d_b_fake)
        tf.summary.scalar('L_D_B_Real', l_d_b_real)
        tf.summary.scalar('L_D_B', l_d_b)
        tf.summary.scalar('L_D', l_d)

        tf.summary.scalar('L_G_A', l_g_a)
        tf.summary.scalar('L_G_B', l_g_b)
        tf.summary.scalar('L_Const_A', l_const_a)
        tf.summary.scalar('L_Const_B', l_const_b)
        tf.summary.scalar('L_G', l_g)

    train_op = tf.group(*[train_op_d_a, train_op_d_b, train_op_g])

    return train_op, global_step, outputs


def input_fn(dataset_a, dataset_b, batch_size=1, num_readers=4, is_training=True):
    provider_a = slim.dataset_data_provider.DatasetDataProvider(
        dataset_a,
        num_readers=num_readers,
        common_queue_capacity=20 * batch_size,
        common_queue_min=10 * batch_size)
    provider_b = slim.dataset_data_provider.DatasetDataProvider(
        dataset_b,
        num_readers=num_readers,
        common_queue_capacity=20 * batch_size,
        common_queue_min=10 * batch_size)
    [image_a, bbox_a] = provider_a.get(['image', 'object/bbox'])
    [image_b, bbox_b] = provider_b.get(['image', 'object/bbox'])

    train_image_size = default_image_size

    def add_channel(image, bbox, padding='ZERO'):
        ymin = bbox[0]
        xmin = bbox[1]
        ymax = bbox[2]
        xmax = bbox[3]

        image_shape = tf.to_float(tf.shape(image))
        height = image_shape[0]
        width = image_shape[1]

        bbox_height = (ymax - ymin) * height
        bbox_width = (xmax - xmin) * width
        channel = tf.ones(tf.to_int32(tf.stack([bbox_height, bbox_width])))
        channel = tf.expand_dims(channel, axis=2)

        pad_top = tf.to_int32(ymin * height)
        pad_left = tf.to_int32(xmin * width)
        height = tf.to_int32(height)
        width = tf.to_int32(width)
        channel = tf.image.pad_to_bounding_box(channel, pad_top, pad_left, height, width)
        # TODO: Decide pad one or zero
        if padding == 'ONE':
            channel = tf.ones_like(channel) - channel

        image = tf.concat([image, channel], axis=2)

        return image

    image_a = tf.image.convert_image_dtype(image_a, dtype=tf.float32)
    image_b = tf.image.convert_image_dtype(image_b, dtype=tf.float32)
    # [Num of boxes, 4] => [4]
    bbox_a = tf.squeeze(bbox_a, axis=0)
{'fraction_non_alphanumeric': 0.12654210085040124, 'fraction_numerical': 0.013654329859863456, 'mean_word_length': 2.4833124739257406, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6528975', 'n_tokens_mistral': 6151, 'n_tokens_neox': 5718, 'n_words': 919}