<jupyter_start><jupyter_text>Utils.py<jupyter_code># Utils code
"""
Some codes from https://github.com/Newmu/dcgan_code
"""
from __future__ import division
import math
import json
import random
import pprint
import scipy.misc
import imageio
import numpy as np
from time import gmtime, strftime
from six.moves import xrange
import PIL
from PIL import Image
# import tensorflow as tf
# # tf.disable_v2_behavior()
# print(tf.__version__)
# import tensorflow.contrib.slim as slim
import tf_slim as slim

pp = pprint.PrettyPrinter()

def show_all_variables():
    model_vars = tf.trainable_variables()
    slim.model_analyzer.analyze_vars(model_vars, print_info=True)


def save_npy_img(images, size, image_path):
    palette=[]
    for i in range(256):
        palette.extend((i,i,i))
    
    palette[:3*21]=np.array([[0, 0, 0],
                            [128, 0, 0],
                            [0, 128, 0],
                            [0, 0, 128],
                            [128, 128, 0],
                            [128, 0, 128],
                            [0, 128, 128],
                            [128, 128, 128],
                            [64, 0, 0]], dtype='uint8').flatten()
    
    cls_map_all = np.zeros((images.shape[0], images.shape[1], images.shape[2], 3), dtype=np.uint8)

    for img_ind in range(images.shape[0]):
        binary_mask = images[img_ind, :, :, :]

        # Add background
        image_sum = np.sum(binary_mask, axis=-1)
        ind = np.where(image_sum==0)
        image_bk = np.zeros((binary_mask.shape[0], binary_mask.shape[1]), dtype=np.float32)
        image_bk[ind] = 1.0
        image_bk = np.reshape(image_bk, (binary_mask.shape[0], binary_mask.shape[1], 1))
        binary_mask = np.concatenate((image_bk, binary_mask), axis=-1)

        cls_map = np.zeros((binary_mask.shape[0], binary_mask.shape[1]), dtype=np.float32)
        cls_map = np.argmax(binary_mask, axis=2)

        cls_map_img = Image.fromarray(cls_map.astype(np.uint8))
        cls_map_img.putpalette(palette)
        cls_map_img = cls_map_img.convert('RGB')
        cls_map_all[img_ind, :, :, :] = np.array(cls_map_img)

    cls_map_all = np.squeeze(merge(cls_map_all, size))
    return imageio.imwrite(image_path, cls_map_all)


def merge(images, size):
    h, w = images.shape[1], images.shape[2]
    if (images.shape[3] in (3,4)):
        c = images.shape[3]
        img = np.zeros((h * size[0], w * size[1], c))
        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
    elif images.shape[3]==1:
        img = np.zeros((h * size[0], w * size[1]))
        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[:,:,0]
        return img
    else:
        raise ValueError('in merge(images,size) images parameter ' 'must have dimensions: HxW or HxWx3 or HxWx4')


def image_manifold_size(num_images):
    manifold_h = int(np.floor(np.sqrt(num_images)))
    manifold_w = int(np.ceil(np.sqrt(num_images)))
    assert manifold_h * manifold_w == num_images
    return manifold_h, manifold_w<jupyter_output><empty_output><jupyter_text>ops.py<jupyter_code>import math
import numpy as np 
# import tensorflow as tf
from tensorflow.python.framework import ops
# from utils import *

class batch_norm(object):
    def __init__(self, epsilon=1e-5, momentum = 0.9, name="batch_norm"):
        with tf.variable_scope(name):
            self.epsilon  = epsilon
            self.momentum = momentum
            self.name = name

    def __call__(self, x, train=True):
        return tf.contrib.layers.batch_norm(x,decay=self.momentum, updates_collections=None, epsilon=self.epsilon, scale=True, is_training=train, scope=self.name)


def conv2d(input_, output_dim, k_h=5, k_w=5, d_h=2, d_w=2, stddev=0.02, padding='SAME', name="conv2d"):
    with tf.variable_scope(name, reuse=tf.AUTO_REUSE):
        if name == 'bbox_pred':
            w = tf.get_variable('w', [k_h, k_w, input_.get_shape()[-1], output_dim], initializer=tf.truncated_normal_initializer(0.0, stddev=0.001))
        else:
            w = tf.get_variable('w', [k_h, k_w, input_.get_shape()[-1], output_dim], initializer=tf.truncated_normal_initializer(stddev=stddev))
        conv = tf.nn.conv2d(input_, w, strides=[1, d_h, d_w, 1], padding=padding)

        biases = tf.get_variable('biases', [output_dim], initializer=tf.constant_initializer(0.0))
        conv = tf.reshape(tf.nn.bias_add(conv, biases), conv.get_shape())

        return conv


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):
    shape = input_.get_shape().as_list()

    with tf.variable_scope(scope or "Linear", reuse=tf.AUTO_REUSE):
        matrix = tf.get_variable("Matrix", [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 relation_nonLocal(input_, name="relation_nonLocal"):
    with tf.variable_scope(name, reuse=tf.AUTO_REUSE):
        shape_org = input_.get_shape().as_list()
        N, H, W, C = shape_org[0], shape_org[1], shape_org[2], shape_org[3]
        # output_dim, d_k, d_g = C/2, C/2, C/2
        output_dim, d_k, d_g = C, C, C

        f_v = conv2d(input_, output_dim, k_h=1, k_w=1, d_h=1, d_w=1, name="f_v")
        f_k = conv2d(input_, d_k, k_h=1, k_w=1, d_h=1, d_w=1, name="f_k")
        f_q = conv2d(input_, d_k, k_h=1, k_w=1, d_h=1, d_w=1, name="f_q")

        f_k = tf.reshape(f_k, [N, H*W, d_k])
        f_q = tf.transpose(tf.reshape(f_q, [N, H*W, d_k]), perm=[0, 2, 1])
        w = tf.matmul(f_k, f_q)/(H*W)

        f_r = tf.matmul(tf.transpose(w, perm=[0, 2, 1]), tf.reshape(f_v, [N, H*W, output_dim]))
        f_r = tf.reshape(f_r, [N, H, W, output_dim])
        f_r = conv2d(f_r, C, k_h=1, k_w=1, d_h=1, d_w=1, name="f_r")

        return f_r 


def layout_point(final_pred, output_height, output_width, name="layout_point"):
    with tf.variable_scope(name, reuse=tf.AUTO_REUSE):

        bbox_pred = tf.reshape(final_pred, [64, 128, 2])

        x_r = tf.reshape(tf.range(output_width, dtype=tf.float32), [1, output_width, 1, 1])
        x_r = tf.reshape(tf.tile(x_r, [1, 1, output_width, 1]), [1, output_width*output_width, 1, 1])
        x_r = tf.tile(x_r, [64, 1, 128, 1])

        y_r = tf.reshape(tf.range(output_height, dtype=tf.float32), [1, 1, output_height, 1])
        y_r = tf.reshape(tf.tile(y_r, [1, output_height, 1, 1]), [1, output_height*output_height, 1, 1])
        y_r = tf.tile(y_r, [64, 1, 128, 1])

        x_pred = tf.reshape(tf.slice(bbox_pred, [0, 0, 0], [-1, -1, 1]), [64, 1, 128, 1])
        x_pred = tf.tile(x_pred, [1, output_width*output_width, 1, 1])
        x_pred = (output_width-1.0) * x_pred

        y_pred = tf.reshape(tf.slice(bbox_pred, [0, 0, 1], [-1, -1, 1]), [64, 1, 128, 1])
        y_pred = tf.tile(y_pred, [1, output_height*output_height, 1, 1])
        y_pred = (output_height-1.0) * y_pred

        x_diff = tf.maximum(0.0, 1.0-tf.abs(x_r - x_pred))
        y_diff = tf.maximum(0.0, 1.0-tf.abs(y_r - y_pred))
        xy_diff = x_diff * y_diff

        xy_max = tf.nn.max_pool(xy_diff, ksize=[1, 1, 128, 1], strides=[1, 1, 1, 1], padding='VALID')
        xy_max = tf.reshape(xy_max, [64, output_height, output_width, 1])

        return xy_max 


# For bbox layout generation 
def layout_bbox(final_pred, output_height, output_width, name="layout_bbox"):
    with tf.variable_scope(name, reuse=tf.AUTO_REUSE):
        final_pred = tf.reshape(final_pred, [64, 9, 9])
        bbox_reg = tf.slice(final_pred, [0, 0, 0], [-1, -1, 4])
        cls_prob = tf.slice(final_pred, [0, 0, 4], [-1, -1, 5])

        bbox_reg = tf.reshape(bbox_reg, [64, 9, 4])

        x_c = tf.slice(bbox_reg, [0, 0, 0], [-1, -1, 1]) * output_width
        y_c = tf.slice(bbox_reg, [0, 0, 1], [-1, -1, 1]) * output_height
        w   = tf.slice(bbox_reg, [0, 0, 2], [-1, -1, 1]) * output_width
        h   = tf.slice(bbox_reg, [0, 0, 3], [-1, -1, 1]) * output_height

        x1 = x_c - 0.5*w
        x2 = x_c + 0.5*w
        y1 = y_c - 0.5*h
        y2 = y_c + 0.5*h

        xt = tf.reshape(tf.range(output_width, dtype=tf.float32), [1, 1, 1, -1])
        xt = tf.reshape(tf.tile(xt, [64, 9, output_height, 1]), [64, 9, -1])

        yt = tf.reshape(tf.range(output_height, dtype=tf.float32), [1, 1, -1, 1])
        yt = tf.reshape(tf.tile(yt, [64, 9, 1, output_width]), [64, 9, -1])

        x1_diff = tf.reshape(xt-x1, [64, 9, output_height, output_width, 1])
        y1_diff = tf.reshape(yt-y1, [64, 9, output_height, output_width, 1])
        x2_diff = tf.reshape(x2-xt, [64, 9, output_height, output_width, 1])
        y2_diff = tf.reshape(y2-yt, [64, 9, output_height, output_width, 1])

        x1_line = tf.nn.relu(1.0 - tf.abs(x1_diff)) * tf.minimum(tf.nn.relu(y1_diff), 1.0) * tf.minimum(tf.nn.relu(y2_diff), 1.0)
        x2_line = tf.nn.relu(1.0 - tf.abs(x2_diff)) * tf.minimum(tf.nn.relu(y1_diff), 1.0) * tf.minimum(tf.nn.relu(y2_diff), 1.0)
        y1_line = tf.nn.relu(1.0 - tf.abs(y1_diff)) * tf.minimum(tf.nn.relu(x1_diff), 1.0) * tf.minimum(tf.nn.relu(x2_diff), 1.0)
        y2_line = tf.nn.relu(1.0 - tf.abs(y2_diff)) * tf.minimum(tf.nn.relu(x1_diff), 1.0) * tf.minimum(tf.nn.relu(x2_diff), 1.0)

        xy_max = tf.reduce_max(tf.concat([x1_line, x2_line, y1_line, y2_line], axis=-1), axis=-1, keep_dims=True)

        spatial_prob = tf.multiply(tf.tile(xy_max, [1, 1, 1, 1, 5]), tf.reshape(cls_prob, [64, 9, 1, 1, 5]))
        spatial_prob_max = tf.reduce_max(spatial_prob, axis=1, keep_dims=False)

        return spatial_prob_max<jupyter_output><empty_output><jupyter_text>model.py<jupyter_code>from __future__ import division
import os
import time
import math
from glob import glob
# import tensorflow as tf
import numpy as np
from six.moves import xrange
import random

# from ops import *
# from utils import *

os.environ["CUDA_VISIBLE_DEVICES"]="1"

class LAYOUTGAN(object):
    def __init__(self, sess, batch_size=64, sample_num=64, dataset_name='default', checkpoint_dir=None, sample_dir=None):
        """
        Args:
            sess: TensorFlow session
            batch_size: The size of batch.
        """
        self.sess = sess
        self.batch_size = batch_size
        self.sample_num = sample_num
        self.dataset_name = dataset_name
        self.checkpoint_dir = checkpoint_dir
        self.d_bn0 = batch_norm(name='d_bn0')
        self.d_bn1 = batch_norm(name='d_bn1')
        self.d_bn2 = batch_norm(name='d_bn2')
        self.g_bn0_0 = batch_norm(name='g_bn0_0')
        self.g_bn0_1 = batch_norm(name='g_bn0_1')
        self.g_bn0_2 = batch_norm(name='g_bn0_2')
        self.g_bn0_3 = batch_norm(name='g_bn0_3')
        self.g_bn1_0 = batch_norm(name='g_bn1_0')
        self.g_bn1_1 = batch_norm(name='g_bn1_1')
        self.g_bn1_2 = batch_norm(name='g_bn1_2')
        self.g_bn1_3 = batch_norm(name='g_bn1_3')
        self.g_bn_x0 = batch_norm(name='g_bn_x0')
        self.g_bn_x1 = batch_norm(name='g_bn_x1')
        self.g_bn_x2 = batch_norm(name='g_bn_x2')
        self.g_bn_x3 = batch_norm(name='g_bn_x3')
        self.g_bn_x4 = batch_norm(name='g_bn_x4')
        self.g_bn_x5 = batch_norm(name='g_bn_x5')
        self.g_bn_x6 = batch_norm(name='g_bn_x6')
        self.g_bn_x7 = batch_norm(name='g_bn_x7')

        # self.data_pre = np.load('./data/pre_data_cls.npy')
        # print("complete loading pre_dat.npy")
        # print(len(self.data_pre))

        # Modified block to ingest data depending on the type of dataset mentioned
        npy_path = './data/' + self.dataset_name + '_train.npy'
        self.data_pre = np.load(npy_path)
        print("complete loading " + npy_path)
        print(len(self.data_pre))

        self.build_model()

    def build_model(self):
        self.inputs = tf.placeholder(tf.float32, [self.batch_size, 9, 9], name='real_images')
        self.z = tf.placeholder(tf.float32, [self.batch_size, 9, 9], name='z')

        self.G = self.generator(self.z)
        self.D, self.D_logits = self.discriminator(self.inputs, reuse=False)
        self.G_sample = self.sampler(self.z)
        self.D_, self.D_logits_ = self.discriminator(self.G, reuse=True)

        def sigmoid_cross_entropy_with_logits(x, y):
            try:
                return tf.nn.sigmoid_cross_entropy_with_logits(logits=x, labels=y)
            except:
                return tf.nn.sigmoid_cross_entropy_with_logits(logits=x, targets=y)

        # Discriminator loss
        self.d_loss_real = tf.reduce_mean(sigmoid_cross_entropy_with_logits(self.D_logits, tf.ones_like(self.D)))
        self.d_loss_fake = tf.reduce_mean(sigmoid_cross_entropy_with_logits(self.D_logits_, tf.zeros_like(self.D_)))
        # normalize_g = tf.nn.l2_normalize(self.G, axis=-1)
        # self.d_batch_similarity = tf.reduce_mean(tf.matmul(normalize_g, tf.transpose(normalize_g, perm=[0,2,1]))) # Penalty on similarity
        self.d_loss = self.d_loss_real + self.d_loss_fake # + self.d_batch_similarity 

        # Generator Loss
        lambda_bbox = 2
        lambda_class = 1
        self.g_loss_wireframe = tf.reduce_mean(sigmoid_cross_entropy_with_logits(self.D_logits_, tf.ones_like(self.D_)))
        self.g_loss_bbox = tf.reduce_mean(tf.squared_difference(
            tf.slice(self.G, (0, 0, 0), (self.batch_size, 9, 4)), 
            tf.slice(self.inputs, (0, 0, 0), (self.batch_size, 9, 4))))
        self.g_loss_class = tf.reduce_mean(sigmoid_cross_entropy_with_logits(
            tf.slice(self.G, (0, 0, 4), (self.batch_size, 9, 5)), 
            tf.slice(self.inputs, (0, 0, 4), (self.batch_size, 9, 5))))
        self.g_loss_reconstruction = (lambda_bbox * self.g_loss_bbox) + (lambda_class * self.g_loss_class)
        self.g_loss = self.g_loss_wireframe # + self.g_loss_reconstruction

        self.d_vars = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, scope="discriminator")
        self.g_vars = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, scope="generator")

        self.saver = tf.train.Saver()
        self.summ = tf.summary.merge([
            tf.summary.scalar('d_loss_real', self.d_loss_real),
            tf.summary.scalar('d_loss_fake', self.d_loss_fake),
            tf.summary.scalar('d_loss', self.d_loss),
            tf.summary.scalar('g_loss', self.g_loss),
            tf.summary.scalar('g_loss_reconstruction', self.g_loss_reconstruction),
        ])

    def train(self, config):
        global_step = tf.Variable(0, trainable=False)
        epoch_step = len(self.data_pre) // config.batch_size    
        lr = tf.train.exponential_decay(0.00001, global_step, 20*epoch_step, 0.1, staircase=True)
        lr_ = tf.summary.scalar('learning_rate', lr)

        d_optim = tf.train.AdamOptimizer(lr, beta1=0.9).minimize(self.d_loss, var_list=self.d_vars, global_step=global_step)
        g_optim = tf.train.AdamOptimizer(lr, beta1=0.9).minimize(self.g_loss, var_list=self.g_vars)
        g_optim_reconstruction = tf.train.AdamOptimizer(lr, beta1=0.9).minimize(self.g_loss_reconstruction, var_list=self.g_vars)

        try:
            tf.global_variables_initializer().run()
        except:
            tf.initialize_all_variables().run()
        
        sample = self.data_pre[0:self.sample_num]
        sample_inputs = np.array(sample).astype(np.float32)
        # sample_inputs = sample_inputs * 28.0 / 27.0 

        # save partial training data
        sample_dir = os.path.join(config.sample_dir, self.model_dir)
        if not os.path.exists(sample_dir):
            os.makedirs(sample_dir)
        samples = self.sess.run(layout_bbox(self.inputs, 60, 40), feed_dict={self.inputs: sample_inputs})
        size = image_manifold_size(samples.shape[0])
        path = './{}/sample.jpg'.format(sample_dir)
        save_npy_img(samples, size, path)

        # sample_z = np.random.normal(0.5, 0.15, (self.batch_size, 128, 2))
        sample_z_bbox = np.random.normal(0.5, 0.15, (self.batch_size, 9, 4))
        # sample_z_cls = np.identity(5)[np.random.randint(5, size=(self.batch_size, 9))]
        sample_z_cls = sample_inputs[:, :, 4:] # Taking original classes of inputs
        sample_z = np.concatenate([sample_z_bbox, sample_z_cls], axis=-1)
    
        counter = 1
        start_time = time.time()
        could_load, checkpoint_counter = self.load(self.checkpoint_dir)
        if could_load:
            counter = checkpoint_counter
            print(" [*] Load SUCCESS")
        else:
            print(" [!] Load failed...")

        for epoch in xrange(config.epoch):
            np.random.shuffle(self.data_pre)
            batch_idxs = len(self.data_pre) // config.batch_size
            for idx in xrange(0, batch_idxs):
                batch = self.data_pre[idx*config.batch_size:(idx+1)*config.batch_size]
                batch_images = np.array(batch).astype(np.float32)
                batch_images = batch_images * 28.0 / 27.0

                # batch_z = np.random.normal(0.5, 0.15, (64, 128, 2))
                batch_z_bbox = np.random.normal(0.5, 0.15, (self.batch_size, 9, 4))
                # batch_z_cls = np.identity(5)[np.random.randint(5, size=(self.batch_size, 9))]
                batch_z_cls = batch_images[:, :, 4:] # Taking original classes of inputs
                batch_z = np.concatenate([batch_z_bbox, batch_z_cls], axis=-1)

                # Update D network
                _ = self.sess.run([d_optim], feed_dict={ self.inputs: batch_images, self.z: batch_z})
                
                # Update G network
                for ae_iter in range(0): # SKIPPING optimization on reconstruction loss
                    _ = self.sess.run([g_optim_reconstruction], feed_dict={ self.inputs: batch_images, self.z: batch_z})
                for layout_iter in range(2):
                    _ = self.sess.run([g_optim], feed_dict={ self.inputs: batch_images, self.z: batch_z})
                
                errD_fake = self.d_loss_fake.eval({ self.z: batch_z})
                errD_real = self.d_loss_real.eval({ self.inputs: batch_images})
                errG = self.g_loss.eval({self.inputs: batch_images, self.z: batch_z})
                errGRecon = self.g_loss_reconstruction.eval({self.inputs: batch_images, self.z: batch_z})

                counter += 1
                if np.mod(counter, 50) == 0:
                    print("Epoch: [%2d] [%4d/%4d] time: %4.4f, lr:%.8f, d_loss: %.4f, g_loss: %.4f, g_loss_reconstruction: %.4f" 
                          % (epoch, idx, batch_idxs, time.time()-start_time, lr.eval(), errD_fake+errD_real, errG, errGRecon))
                    
                if np.mod(counter, 500) == 1:
                    samples, d_loss, g_loss, g_loss_reconstruction = self.sess.run([self.G_sample, self.d_loss, self.g_loss, self.g_loss_reconstruction], 
                                                            feed_dict={self.z: sample_z, self.inputs: sample_inputs})
                    size = image_manifold_size(samples.shape[0])
                    path = './{}/train_{:02d}_{:04d}.jpg'.format(sample_dir, epoch, idx)
                    save_npy_img(samples, size, path)
                    print("[Sample] d_loss: %.8f, g_loss: %.8f, g_loss_reconstruction: %.8f" % (d_loss, g_loss, g_loss_reconstruction)) 
                
                if np.mod(counter, 2000) == 0:
                    self.save(config.checkpoint_dir, counter)


    def discriminator(self, image, reuse=False):
        with tf.variable_scope("discriminator") as scope:
            if reuse:
                scope.reuse_variables()

            # layout = layout_point(image, 28, 28, name='layout')
            # For bbox layout generation
            layout = layout_bbox(image, 60, 40, name='layout')

            net = lrelu(self.d_bn0(conv2d(layout, 32, k_h=5, k_w=5, d_h=2, d_w=2, padding='VALID', name='conv1')))
            net = lrelu(self.d_bn1(conv2d(net, 64, k_h=5, k_w=5, d_h=2, d_w=2, padding='VALID', name='conv2')))
            net = tf.reshape(net, [self.batch_size, -1])      
            net = lrelu(self.d_bn2(linear(net, 512, scope='fc2')))
            net = linear(net, 1, 'fc3')

        return tf.nn.sigmoid(net), net


    def generator(self, z):
        with tf.variable_scope("generator") as scope:
            # gnet = tf.reshape(z, [64, 128, 1, 2])
            # h0_0 = self.g_bn0_0(conv2d(gnet, 1024, k_h=1, k_w=1, d_h=1, d_w=1, name='g_h0_0'))
            # h0_1 = tf.nn.relu(self.g_bn0_1(conv2d(gnet, 256, k_h=1, k_w=1, d_h=1, d_w=1, name='g_h0_1')))
            # h0_2 = tf.nn.relu(self.g_bn0_2(conv2d(h0_1, 256, k_h=1, k_w=1, d_h=1, d_w=1, name='g_h0_2')))
            # h0_3 = self.g_bn0_3(conv2d(h0_2, 1024, k_h=1, k_w=1, d_h=1, d_w=1, name='g_h0_3'))
            # gnet = tf.nn.relu(tf.add(h0_0, h0_3))

            # For bbox layout generation
            # gnet = tf.reshape(z, [64, 9, 6, 4])
            gnet = tf.reshape(z, [self.batch_size, 9, 1, 4 + 5])
            h0_0 = self.g_bn0_0(conv2d(gnet, 256, k_h=1, k_w=1, d_h=1, d_w=1, name='g_h0_0'))
            h0_1 = tf.nn.relu(self.g_bn0_1(conv2d(gnet, 64, k_h=1, k_w=1, d_h=1, d_w=1, name='g_h0_1')))
            h0_2 = tf.nn.relu(self.g_bn0_2(conv2d(h0_1, 64, k_h=1, k_w=1, d_h=1, d_w=1, name='g_h0_2')))
            h0_3 = self.g_bn0_3(conv2d(h0_2, 256, k_h=1, k_w=1, d_h=1, d_w=1, name='g_h0_3'))
            gnet = tf.nn.relu(tf.add(h0_0, h0_3))
            gnet = tf.reshape(gnet, [self.batch_size, 9, 1, 256])

            # gnet = tf.reshape(gnet, [64, 128, 1, 1024])
            gnet = tf.nn.relu(self.g_bn_x1( tf.add(gnet, self.g_bn_x0(relation_nonLocal(gnet, name='g_non0')))))
            gnet = tf.nn.relu(self.g_bn_x3( tf.add(gnet, self.g_bn_x2(relation_nonLocal(gnet, name='g_non2')))))

            h1_0 = self.g_bn1_0(conv2d(gnet, 1024, k_h=1, k_w=1, d_h=1, d_w=1, name='g_h1_0'))
            h1_1 = tf.nn.relu(self.g_bn1_1(conv2d(h1_0, 256, k_h=1, k_w=1, d_h=1, d_w=1, name='g_h1_1')))
            h1_2 = tf.nn.relu(self.g_bn1_2(conv2d(h1_1, 256, k_h=1, k_w=1, d_h=1, d_w=1, name='g_h1_2')))
            h1_3 = self.g_bn1_3(conv2d(h1_2, 1024, k_h=1, k_w=1, d_h=1, d_w=1, name='g_h1_3'))
            gnet = tf.nn.relu(tf.add(h1_0, h1_3))

            # For bbox layout generation
            # May add more self-attention refinement steps
            gnet = tf.nn.relu(self.g_bn_x5( tf.add(gnet, self.g_bn_x4(relation_nonLocal(gnet, name='g_non4')))))
            gnet = tf.nn.relu(self.g_bn_x7( tf.add(gnet, self.g_bn_x6(relation_nonLocal(gnet, name='g_non6')))))


            # ----  From original repo ---- #

            # # bbox_pred = conv2d(gnet, 2, k_h=1, k_w=1, d_h=1, d_w=1, name='bbox_pred')
            # # bbox_pred = tf.sigmoid(tf.reshape(bbox_pred, [-1, 128, 2]))
            # # final_pred = bbox_pred

            # # For bbox layout generation 
            # cls_score = conv2d(gnet, 6, k_h=1, k_w=1, d_h=1, d_w=1, name='cls_score')
            # cls_prob  = tf.sigmoid(tf.reshape(cls_score, [-1, 9, 6]))
            # final_pred = tf.concat([bbox_pred, cls_prob], axis=-1)

            # ----  Ends ---- #

            bbox_pred = conv2d(gnet, 4, k_h=1, k_w=1, d_h=1, d_w=1, name='bbox_pred')
            bbox_pred = tf.sigmoid(tf.reshape(bbox_pred, [-1, 9, 4]))
            # final_pred = bbox_pred
            
            # For bbox layout generation 
            cls_score = conv2d(gnet, 5, k_h=1, k_w=1, d_h=1, d_w=1, name='cls_score')
            cls_prob  = tf.sigmoid(tf.reshape(cls_score, [-1, 9, 5]))
            final_pred = tf.concat([bbox_pred, cls_prob], axis=-1)

            return final_pred 


    def sampler(self, z):
        with tf.variable_scope("generator", reuse=tf.AUTO_REUSE) as scope:
            scope.reuse_variables()
            gnet = tf.reshape(z, [self.batch_size, 9, 1, 4 + 5])
            h0_0 = self.g_bn0_0(conv2d(gnet, 256, k_h=1, k_w=1, d_h=1, d_w=1, name='g_h0_0'), train=False)
            h0_1 = tf.nn.relu(self.g_bn0_1(conv2d(gnet, 64, k_h=1, k_w=1, d_h=1, d_w=1, name='g_h0_1'), train=False))
            h0_2 = tf.nn.relu(self.g_bn0_2(conv2d(h0_1, 64, k_h=1, k_w=1, d_h=1, d_w=1, name='g_h0_2'), train=False))
            h0_3 = self.g_bn0_3(conv2d(h0_2, 256, k_h=1, k_w=1, d_h=1, d_w=1, name='g_h0_3'), train=False)
            gnet = tf.nn.relu(tf.add(h0_0, h0_3))
            gnet = tf.reshape(gnet, [self.batch_size, 9, 1, 256])

            # gnet = tf.reshape(gnet, [self.batch_size, 128, 1, 1024])
            gnet = tf.nn.relu(self.g_bn_x1( tf.add(gnet, self.g_bn_x0(relation_nonLocal(gnet, name='g_non0'), train=False)), train=False))
            gnet = tf.nn.relu(self.g_bn_x3( tf.add(gnet, self.g_bn_x2(relation_nonLocal(gnet, name='g_non2'), train=False)), train=False))

            h1_0 = self.g_bn1_0(conv2d(gnet, 1024, k_h=1, k_w=1, d_h=1, d_w=1, name='g_h1_0'), train=False)
            h1_1 = tf.nn.relu(self.g_bn1_1(conv2d(h1_0, 256, k_h=1, k_w=1, d_h=1, d_w=1, name='g_h1_1'), train=False))
            h1_2 = tf.nn.relu(self.g_bn1_2(conv2d(h1_1, 256, k_h=1, k_w=1, d_h=1, d_w=1, name='g_h1_2'), train=False))
            h1_3 = self.g_bn1_3(conv2d(h1_2, 1024, k_h=1, k_w=1, d_h=1, d_w=1, name='g_h1_3'), train=False)
            gnet = tf.nn.relu(tf.add(h1_0, h1_3))

            gnet = tf.nn.relu(self.g_bn_x5( tf.add(gnet, self.g_bn_x4(relation_nonLocal(gnet, name='g_non4'), train=False)), train=False))
            gnet = tf.nn.relu(self.g_bn_x7( tf.add(gnet, self.g_bn_x6(relation_nonLocal(gnet, name='g_non6'), train=False)), train=False))

            bbox_pred = conv2d(gnet, 4, k_h=1, k_w=1, d_h=1, d_w=1, name='bbox_pred')
            bbox_pred = tf.sigmoid(tf.reshape(bbox_pred, [-1, 9, 4]))
            # final_pred = bbox_pred

            cls_score = conv2d(gnet, 5, k_h=1, k_w=1, d_h=1, d_w=1, name='cls_score')
            cls_prob  = tf.sigmoid(tf.reshape(cls_score, [-1, 9, 5]))
            final_pred = tf.concat([bbox_pred, cls_prob], axis=-1)

            layout = layout_bbox(final_pred, 60, 40, name='layout')

            return layout


    @property
    def model_dir(self):
        return "{}_{}".format(self.dataset_name, self.batch_size)
      
    def save(self, checkpoint_dir, step):
        model_name = "LAYOUTGAN.model"
        checkpoint_dir = os.path.join(checkpoint_dir, self.model_dir)

        if not os.path.exists(checkpoint_dir):
            os.makedirs(checkpoint_dir)
        
        self.saver.save(self.sess, os.path.join(checkpoint_dir, model_name), global_step=step)

    def load(self, checkpoint_dir):
        import re
        print(" [*] Reading checkpoints...")
        checkpoint_dir = os.path.join(checkpoint_dir, self.model_dir)

        ckpt = tf.train.get_checkpoint_state(checkpoint_dir)
        if ckpt and ckpt.model_checkpoint_path:
            ckpt_name = os.path.basename(ckpt.model_checkpoint_path)
            self.saver.restore(self.sess, os.path.join(checkpoint_dir, ckpt_name))
            counter = int(next(re.finditer("(\d+)(?!.*\d)",ckpt_name)).group(0))
            print(" [*] Success to read {}".format(ckpt_name))
            return True, counter
        else:
            print(" [*] Failed to find a checkpoint")
            return False, 0<jupyter_output><empty_output><jupyter_text>main.py<jupyter_code>import os
import scipy.misc
import imageio
import numpy as np

# from model import LAYOUTGAN
# from utils import pp, show_all_variables

def del_all_flags(FLAGS):
    keys_list = ["epoch", "batch_size", "dataset", "checkpoint_dir", "sample_dir", "train"]   
    for key in keys_list:
        if key in FLAGS._flags():
            FLAGS.__delattr__(key)

flags = tf.app.flags
del_all_flags(flags.FLAGS)
flags.DEFINE_integer("epoch", 50, "Epoch to train [25]")
flags.DEFINE_integer("batch_size", 64, "The size of batch images [64]")
flags.DEFINE_string("dataset", "doc", "The name of dataset [mnist, doc, clipart, tangram]")
flags.DEFINE_string("checkpoint_dir", "vanilla_checkpoint", "Directory name to save the checkpoints [checkpoint]")
flags.DEFINE_string("sample_dir", "vanilla_samples", "Directory name to save the image samples [samples]")
flags.DEFINE_boolean("train", True, "True for training, False for testing [False]")
FLAGS = flags.FLAGS
            
def main(_):
    pp.pprint(flags.FLAGS.__flags)
    if not os.path.exists(FLAGS.checkpoint_dir):
        os.makedirs(FLAGS.checkpoint_dir)
    if not os.path.exists(FLAGS.sample_dir):
        os.makedirs(FLAGS.sample_dir)
    
    gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.333)
    run_config = tf.ConfigProto()
    run_config.gpu_options.allow_growth=True
    with tf.Session(config=run_config) as sess:
        layoutgan = LAYOUTGAN(sess, batch_size=FLAGS.batch_size, sample_num=FLAGS.batch_size, dataset_name=FLAGS.dataset, checkpoint_dir=FLAGS.checkpoint_dir, sample_dir=FLAGS.sample_dir)
    
        show_all_variables()
        if FLAGS.train:
            print("\n ------- \n Training ... \n")
            layoutgan.train(FLAGS)
        else:
            print("\n ------- \n Loading checkpoint ... \n")
            if not layoutgan.load(FLAGS.checkpoint_dir)[0]:
                raise Exception("[!] Train a model first, then run test mode")

if __name__ == '__main__':
    tf.app.run()<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.14928148774302621, 'fraction_numerical': 0.03928994082840237, 'mean_word_length': 2.7174459527400705, 'pattern_counts': {'":': 0, '<': 17, '<?xml version=': 0, '>': 17, 'https://': 1, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '4', 'dir': 'jupyter-structured-clean-dedup', 'id': '160369', 'license': 'Apache-2.0', 'n_tokens_mistral': 12182, 'n_tokens_neox': 11517, 'n_words': 2059, 'path': 'Object_Alignment/LayoutGAN-Document-Vanilla.ipynb', 'repo_name': 'hassan-mahmood/Layout-Agnostic-Object-Alignment-and-Image-Generation', 'size': '29575'}