[ { "path": "LICENSE", "content": "MIT License\n\nCopyright (c) 2021 Layout-Generation\n\nPermission is hereby granted, free of charge, to any person obtaining a copy\nof this software and associated documentation files (the \"Software\"), to deal\nin the Software without restriction, including without limitation the rights\nto use, copy, modify, merge, publish, distribute, sublicense, and/or sell\ncopies of the Software, and to permit persons to whom the Software is\nfurnished to do so, subject to the following conditions:\n\nThe above copyright notice and this permission notice shall be included in all\ncopies or substantial portions of the Software.\n\nTHE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR\nIMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,\nFITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE\nAUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER\nLIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,\nOUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE\nSOFTWARE.\n" }, { "path": "Layout Transformer/Notebook/Data/.gitkeep", "content": "" }, { "path": "Layout Transformer/Notebook/Layout_Transformer.ipynb", "content": "{\n \"cells\": [\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {\n \"id\": \"NSLt1-42rWxE\"\n },\n \"source\": [\n \"### **Linking Storage Drive**\\n\",\n \"---\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"metadata\": {\n \"colab\": {\n \"base_uri\": \"https://localhost:8080/\"\n },\n \"id\": \"lD3wobXKqr-f\",\n \"outputId\": \"fb5730a7-516a-456e-bfb4-9a6977426992\"\n },\n \"outputs\": [],\n \"source\": [\n \"from google.colab import drive\\n\",\n \"drive.mount('/content/drive')\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"metadata\": {\n \"id\": \"-B6ea86VOmGK\"\n },\n \"outputs\": [],\n \"source\": [\n \"root = '/LayoutTransformer/'\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {\n \"id\": \"_WC6MGpirp9N\"\n },\n \"source\": [\n \"### **Imports**\\n\",\n \"---\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"metadata\": {\n \"id\": \"-dcOiEoOrVqq\"\n },\n \"outputs\": [],\n \"source\": [\n \"import tensorflow as tf\\n\",\n \"from tensorflow.keras.layers import Layer\\n\",\n \"from tensorflow.keras import Model\\n\",\n \"from tensorflow.keras import backend as k\\n\",\n \"\\n\",\n \"import numpy as np\\n\",\n \"import json\\n\",\n \"import os\\n\",\n \"import gc\\n\",\n \"\\n\",\n \"import matplotlib.pyplot as plt\\n\",\n \"from matplotlib import gridspec\\n\",\n \"from matplotlib.patches import Patch\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {\n \"id\": \"Jp4KLhN7J23d\"\n },\n \"source\": [\n \"### **GPU Setup**\\n\",\n \"---\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"metadata\": {\n \"id\": \"fcMPzjTTJ3At\"\n },\n \"outputs\": [],\n \"source\": [\n \"os.environ['CUDA_VISIBLE_DEVICES']='0'\\n\",\n \"config = tf.compat.v1.ConfigProto()\\n\",\n \"config.gpu_options.per_process_gpu_memory_fraction = 0.9\\n\",\n \"session = tf.compat.v1.Session(config=config)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {\n \"id\": \"nrxAk-K7sPwB\"\n },\n \"source\": [\n \"### **Importing Data**\\n\",\n \"\\n\",\n \"Data format:\\n\",\n \" [ Number of samples x Number of Boxes x [Class,X,Y,W,H] ]\\n\",\n \"\\n\",\n \"* PublayNet\\n\",\n \"* Rico\\n\",\n \"---\\n\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"metadata\": {\n \"id\": \"ghN8p4iQrplX\"\n },\n \"outputs\": [],\n \"source\": [\n \"publaynet_data = np.load(root+'Data/publaynet.npy')\\n\",\n \"rico_data = np.load(root+'Data/rico_new.npy')\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {\n \"id\": \"udEOmzv2HdGN\"\n },\n \"source\": [\n \"### **Layers**\\n\",\n \"\\n\",\n \"---\\n\",\n \"\\n\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"metadata\": {\n \"id\": \"epBvDnI5HdO6\"\n },\n \"outputs\": [],\n \"source\": [\n \"class MMHSALayer(Layer):\\n\",\n \" '''\\n\",\n \" **Masked Multiheaded Self Attention Layer**\\n\",\n \"\\n\",\n \" heads : Specify the number of heads\\n\",\n \" '''\\n\",\n \" def __init__(self,heads=8):\\n\",\n \" super(MMHSALayer, self).__init__()\\n\",\n \" self.heads = heads\\n\",\n \"\\n\",\n \" def build(self,input_shape):\\n\",\n \" self.model_dim = input_shape[-2]\\n\",\n \" self.k = self.add_weight(shape=(self.heads,self.model_dim,self.model_dim),\\n\",\n \" initializer='random_normal',\\n\",\n \" trainable=True,\\n\",\n \" name=\\\"Key\\\")\\n\",\n \" self.q = self.add_weight(shape=(self.heads,self.model_dim,self.model_dim),\\n\",\n \" initializer='random_normal',\\n\",\n \" trainable=True,\\n\",\n \" name=\\\"Query\\\")\\n\",\n \" self.v = self.add_weight(shape=(self.heads,self.model_dim,self.model_dim),\\n\",\n \" initializer='random_normal',\\n\",\n \" trainable=True,\\n\",\n \" name=\\\"Value\\\")\\n\",\n \" self.o = self.add_weight(shape=(self.model_dim,self.model_dim*self.heads),\\n\",\n \" initializer='random_normal',\\n\",\n \" trainable=True,\\n\",\n \" name=\\\"Heads\\\")\\n\",\n \"\\n\",\n \" def call(self,inputs):\\n\",\n \" mask_shape = inputs.shape[-1] \\n\",\n \"\\n\",\n \" mask_0 = np.ones((mask_shape,mask_shape))\\n\",\n \" for i in range(mask_shape):\\n\",\n \" for j in range(mask_shape):\\n\",\n \" if (i>j):\\n\",\n \" mask_0[i][j]=0\\n\",\n \" self.mask_0 = tf.constant(mask_0,dtype=tf.float32)\\n\",\n \"\\n\",\n \" mask_inf = np.zeros((mask_shape,mask_shape))\\n\",\n \" for i in range(mask_shape):\\n\",\n \" for j in range(mask_shape):\\n\",\n \" if (i>j):\\n\",\n \" mask_inf[i][j]=-10000000000\\n\",\n \" self.mask_inf = tf.constant(mask_inf,dtype=tf.float32)\\n\",\n \"\\n\",\n \" inputs = tf.expand_dims(inputs,1)\\n\",\n \"\\n\",\n \" key=tf.matmul(self.k,inputs)\\n\",\n \" que=tf.matmul(self.q,inputs)\\n\",\n \" val=tf.matmul(self.v,inputs)\\n\",\n \"\\n\",\n \" Z=tf.matmul(tf.transpose(key,perm=[0,1,3,2]),que)*(1/np.sqrt(self.model_dim))\\n\",\n \" W=tf.multiply(Z,self.mask_0)\\n\",\n \" W=tf.add(W,self.mask_inf)\\n\",\n \" W=tf.keras.activations.softmax(W,axis=1)\\n\",\n \" W=tf.multiply(W,self.mask_0)\\n\",\n \" W=tf.matmul(val,W)\\n\",\n \"\\n\",\n \" W = tf.reshape(W,(inputs.shape[0],self.model_dim*self.heads,mask_shape))\\n\",\n \"\\n\",\n \" ans = W\\n\",\n \"\\n\",\n \" ans = tf.matmul(self.o,ans)\\n\",\n \" ans=tf.expand_dims(ans,0)\\n\",\n \"\\n\",\n \" ans = tf.squeeze(ans,axis=0)\\n\",\n \"\\n\",\n \" return ans\\n\",\n \"\\n\",\n \"class Dense2D(Layer):\\n\",\n \" '''\\n\",\n \" **2-Dimensional Dense Layer**\\n\",\n \" Applies dense layer column-wise (shared weights). Returns the column size of units.\\n\",\n \"\\n\",\n \" units : Specify the number of output units (column length)\\n\",\n \" '''\\n\",\n \" def __init__(self,units):\\n\",\n \" super(Dense2D, self).__init__()\\n\",\n \" self.units = units\\n\",\n \"\\n\",\n \" def build(self,input_shape):\\n\",\n \" input_len = input_shape[-2]\\n\",\n \"\\n\",\n \" self.w = self.add_weight(shape=(self.units,input_len),\\n\",\n \" initializer='random_normal',\\n\",\n \" trainable=True,\\n\",\n \" name=\\\"dense2dw\\\")\\n\",\n \"\\n\",\n \" def call(self,inputs,activation = None):\\n\",\n \"\\n\",\n \" ans = tf.matmul(self.w,inputs)\\n\",\n \"\\n\",\n \" return ans\\n\",\n \"\\n\",\n \"class FFLayer(Layer):\\n\",\n \" '''\\n\",\n \" **Feed Forward Layer**\\n\",\n \" Applies dense layer column-wise (shared weights), followed by a ReLU Layer, followed by another dense layer column-wise (shared weights). Returns the same column size.\\n\",\n \"\\n\",\n \" dff : Specify the number of units (column length) in the middle layer \\n\",\n \" dropout : Dropout Rate\\n\",\n \" '''\\n\",\n \" def __init__(self, dff=2048, dropout=0.1):\\n\",\n \" super(FFLayer,self).__init__()\\n\",\n \" self.dff = dff \\n\",\n \" self.dropout = dropout \\n\",\n \"\\n\",\n \" def build(self,input_shape):\\n\",\n \" self.dropout = tf.keras.layers.Dropout(self.dropout)\\n\",\n \" self.dout = input_shape[-2]\\n\",\n \"\\n\",\n \" self.w1 = self.add_weight(shape=(self.dff,self.dout),\\n\",\n \" initializer='random_normal',\\n\",\n \" trainable=True,\\n\",\n \" name=\\\"ffw1\\\")\\n\",\n \" self.w2 = self.add_weight(shape=(self.dout,self.dff),\\n\",\n \" initializer='random_normal',\\n\",\n \" trainable=True,\\n\",\n \" name=\\\"ffw2\\\")\\n\",\n \" self.b1 = self.add_weight(shape=(self.dff,1),\\n\",\n \" initializer='random_normal',\\n\",\n \" trainable=True,\\n\",\n \" name=\\\"ffb1\\\")\\n\",\n \" self.b2 = self.add_weight(shape=(self.dout,1),\\n\",\n \" initializer='random_normal',\\n\",\n \" trainable=True,\\n\",\n \" name=\\\"ffb2\\\")\\n\",\n \"\\n\",\n \" def call(self,inputs):\\n\",\n \"\\n\",\n \" ans = tf.add(tf.matmul(self.w1,inputs),self.b1)\\n\",\n \" ans = tf.keras.activations.relu(ans)\\n\",\n \" ans = tf.add(tf.matmul(self.w2,ans),self.b2)\\n\",\n \"\\n\",\n \" ans = self.dropout(ans)\\n\",\n \" \\n\",\n \" return ans\\n\",\n \"\\n\",\n \"class ANLayer(Layer):\\n\",\n \" '''\\n\",\n \" **Add and Normalize Layer**\\n\",\n \" Adds and then Normalizes column wise.\\n\",\n \" '''\\n\",\n \" def __init__(self):\\n\",\n \" super(ANLayer,self).__init__()\\n\",\n \" self.Normal = tf.keras.layers.LayerNormalization(axis=1)\\n\",\n \"\\n\",\n \" def call(self,inputs1,inputs2):\\n\",\n \" sum = tf.add(inputs1,inputs2)\\n\",\n \" ans=self.Normal(sum)\\n\",\n \" return ans\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {\n \"id\": \"qJrGCXu4suIp\"\n },\n \"source\": [\n \"### **Model**\\n\",\n \"---\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"metadata\": {\n \"id\": \"g6cLgJ9tsuPH\"\n },\n \"outputs\": [],\n \"source\": [\n \"class LTModel(Model):\\n\",\n \" def __init__(self, input_shape, layers, heads, dff, model_dim, dropout):\\n\",\n \" super(LTModel, self).__init__()\\n\",\n \"\\n\",\n \" self.emb = Dense2D(model_dim)\\n\",\n \"\\n\",\n \" self.SA = []\\n\",\n \" self.AN1 = []\\n\",\n \" self.FF = []\\n\",\n \" self.AN2 = []\\n\",\n \"\\n\",\n \" for i in range(layers):\\n\",\n \" self.SA.append(MMHSALayer(heads))\\n\",\n \" self.AN1.append(ANLayer())\\n\",\n \" self.FF.append(FFLayer(dff, dropout))\\n\",\n \" self.AN2.append(ANLayer())\\n\",\n \"\\n\",\n \" self.deemb = Dense2D(input_shape)\\n\",\n \" self.sm = tf.keras.layers.Softmax(axis=1)\\n\",\n \" \\n\",\n \" def call(self, x):\\n\",\n \" x = self.emb(x)\\n\",\n \"\\n\",\n \" for i in range(len(self.SA)):\\n\",\n \" y = self.SA[i](x)\\n\",\n \" x = self.AN1[i](x,y)\\n\",\n \" y = self.FF[i](x)\\n\",\n \" x = self.AN2[i](x,y)\\n\",\n \"\\n\",\n \" x = self.deemb(x)\\n\",\n \" x = self.sm(x)\\n\",\n \"\\n\",\n \" return x\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {\n \"id\": \"XkuubrSpswZ7\"\n },\n \"source\": [\n \"### **Layout Transformer**\\n\",\n \"---\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"metadata\": {\n \"id\": \"EEWk72cFzmwX\"\n },\n \"outputs\": [],\n \"source\": [\n \"class LayoutTransformer:\\n\",\n \"\\n\",\n \" def __init__(self, n_classes, class_labels=None, n_anchors=(32,32), d=512, n_layers=6, n_heads=8, dff=2048, dropout=0.1):\\n\",\n \" self.n_classes = n_classes+2\\n\",\n \" self.n_anchors = n_anchors\\n\",\n \" self.d = d\\n\",\n \" self.n_layers = n_layers\\n\",\n \" self.n_heads = n_heads\\n\",\n \" self.dff = dff\\n\",\n \" self.dropout = dropout\\n\",\n \" self.n_row = n_anchors[0]\\n\",\n \" self.n_col = n_anchors[1]\\n\",\n \" self.input_dim = 2+n_classes+2*(n_anchors[0]+n_anchors[1])\\n\",\n \" self.model = LTModel(self.input_dim, model_dim=d, layers=n_layers, heads=n_heads, dff=dff, dropout=dropout)\\n\",\n \" self.loss_his = []\\n\",\n \" self.lr_his = []\\n\",\n \" self.train_data_his = []\\n\",\n \" if class_labels == None:\\n\",\n \" self.labels = range(1,n_classes+1)\\n\",\n \" else:\\n\",\n \" self.labels = class_labels\\n\",\n \" \\n\",\n \" \\n\",\n \" def compile(self, lr=1e-5):\\n\",\n \" self.model.compile(loss=tf.keras.losses.KLDivergence(),\\n\",\n \" metrics = [tf.keras.losses.KLDivergence()],\\n\",\n \" optimizer = tf.keras.optimizers.Adam(learning_rate=lr))\\n\",\n \" \\n\",\n \" \\n\",\n \" def build(self):\\n\",\n \" self.model.build((1,self.input_dim,1))\\n\",\n \" \\n\",\n \"\\n\",\n \" def summary(self):\\n\",\n \" self.build()\\n\",\n \" print(self.model.summary())\\n\",\n \" \\n\",\n \"\\n\",\n \" def train(self, epochs, batch_size=1, train_data_index=\\\"All\\\", rlrop_factor=0.5, rlrop_patience=1000, rlrop_min_delta=0.001):\\n\",\n \" if train_data_index == \\\"All\\\":\\n\",\n \" train_data_index = range(self.data.shape[0])\\n\",\n \" rlrop = tf.keras.callbacks.ReduceLROnPlateau(factor=rlrop_factor,patience=rlrop_patience,verbose=1,min_delta=rlrop_min_delta,monitor='kl_divergence')\\n\",\n \" callbacks = [rlrop]\\n\",\n \" history = self.model.fit(x=tf.convert_to_tensor(self.x_data[train_data_index]), y=tf.convert_to_tensor(self.y_data[train_data_index]), epochs=epochs, batch_size=batch_size, callbacks=callbacks)\\n\",\n \" self.loss_his.extend(history.history['loss'])\\n\",\n \" self.lr_his.extend(history.history['lr'])\\n\",\n \" for i in range(epochs):\\n\",\n \" self.train_data_his.append(len(train_data_index))\\n\",\n \" \\n\",\n \"\\n\",\n \" def load_weights(self, folder_path, filename):\\n\",\n \" self.build()\\n\",\n \" self.model.load_weights(folder_path + '/' + str(filename) + '.h5')\\n\",\n \" his = json.loads(open(folder_path + '/' + str(filename) + '.json').read())\\n\",\n \"\\n\",\n \" self.loss_his = his['loss']\\n\",\n \" self.train_data_his = his['data']\\n\",\n \" self.lr_his = his['lr']\\n\",\n \" \\n\",\n \"\\n\",\n \" def save_weights(self, folder_path, filename):\\n\",\n \" his = json.dumps({'loss':list(np.array(self.loss_his,dtype='float')),'data':list(np.array(self.train_data_his,dtype='float')),'lr':list(np.array(self.lr_his,dtype='float'))})\\n\",\n \" open(folder_path + '/' + str(filename) + '.json','w').write(his)\\n\",\n \" self.model.save_weights(folder_path + '/' + str(filename) + '.h5')\\n\",\n \" \\n\",\n \"\\n\",\n \" def predict(self, input):\\n\",\n \" input = tf.convert_to_tensor(input, dtype='float32')\\n\",\n \" return self.model(input).numpy()\\n\",\n \" \\n\",\n \"\\n\",\n \" def load_data(self, data, rows, cols, e=0.1):\\n\",\n \" # Make number of boxes equal in eeach document\\n\",\n \" max_box = 0\\n\",\n \" for doc in data:\\n\",\n \" max_box = max(max_box,len(doc))\\n\",\n \" for doc in range(len(data)):\\n\",\n \" while (len(data[doc]) and y_data with \\n\",\n \" x_data = []\\n\",\n \" y_data = []\\n\",\n \"\\n\",\n \" for doc in onehot_data:\\n\",\n \" bos = list(np.zeros(self.input_dim))\\n\",\n \" bos[0]=1\\n\",\n \" x = [bos,*doc]\\n\",\n \" x = np.array(x).T\\n\",\n \" x_data.append(x)\\n\",\n \"\\n\",\n \" eos = list(np.zeros(self.input_dim))\\n\",\n \" eos[self.n_classes-1] = 1\\n\",\n \" y = [*doc,eos]\\n\",\n \" for box in y:\\n\",\n \" for k in range(0, self.n_classes):\\n\",\n \" box[k] = (1 - e) * box[k] + e / self.n_classes\\n\",\n \" for k in range(self.n_classes, self.n_classes+self.n_col):\\n\",\n \" box[k] = (1 - e) * box[k] + e / self.n_col\\n\",\n \" for k in range(self.n_classes+self.n_col, self.n_classes+self.n_col+self.n_row):\\n\",\n \" box[k] = (1 - e) * box[k] + e / self.n_row\\n\",\n \" for k in range(self.n_classes+self.n_col+self.n_row, self.n_classes+2*self.n_col+self.n_row):\\n\",\n \" box[k] = (1 - e) * box[k] + e / self.n_col\\n\",\n \" for k in range(self.n_classes+2*self.n_col+self.n_row, self.n_classes+2*self.n_col+2*self.n_row):\\n\",\n \" box[k] = (1 - e) * box[k] + e / self.n_row\\n\",\n \" y = np.array(y).T\\n\",\n \" y_data.append(y)\\n\",\n \"\\n\",\n \" self.x_data = np.array(x_data,dtype=\\\"float32\\\")\\n\",\n \" self.y_data = np.array(y_data,dtype=\\\"float32\\\")\\n\",\n \" \\n\",\n \"\\n\",\n \" def onehot(self,box, prob=[-1,-1,-1,-1,-1]):\\n\",\n \" p = 1\\n\",\n \" c = np.argsort(box[0:self.n_classes],axis=0)[prob[0]]\\n\",\n \" x = np.argsort(box[self.n_classes:self.n_classes+self.n_col],axis=0)[prob[1]] + self.n_classes\\n\",\n \" y = np.argsort(box[self.n_classes+self.n_col:self.n_classes+self.n_col+self.n_row],axis=0)[prob[2]] + self.n_classes+self.n_col\\n\",\n \" w = np.argsort(box[self.n_classes+self.n_col+self.n_row:self.n_classes+2*self.n_col+self.n_row],axis=0)[prob[3]] + self.n_classes+self.n_col+self.n_row\\n\",\n \" h = np.argsort(box[self.n_classes+2*self.n_col+self.n_row:self.n_classes+2*self.n_col+2*self.n_row],axis=0)[prob[4]] + self.n_classes+2*self.n_col+self.n_row\\n\",\n \" p = p*box[c][0]*box[x][0]*box[y][0]*box[w][0]*box[h][0]\\n\",\n \" res = np.zeros((self.input_dim,1))\\n\",\n \" res[c,0]=1\\n\",\n \" res[x,0]=1\\n\",\n \" res[y,0]=1\\n\",\n \" res[w,0]=1\\n\",\n \" res[h,0]=1\\n\",\n \" return (res,p)\\n\",\n \" \\n\",\n \"\\n\",\n \" def sort_prob(self,docs):\\n\",\n \" p = [int(p) for box,p in docs]\\n\",\n \" p = np.argsort(p)[::-1]\\n\",\n \" res = []\\n\",\n \" for i in p:\\n\",\n \" res.append(docs[i])\\n\",\n \" return res\\n\",\n \" \\n\",\n \"\\n\",\n \" def get_color(self,c):\\n\",\n \" color_key = [\\\"#00ffff\\\",\\\"#fff5ee\\\",\\\"#dc143c\\\",\\\"#ffff00\\\",\\\"#00ff00\\\",\\\"#ff00ff\\\",\\\"#1e90ff\\\",\\n\",\n \" \\\"#ff1493\\\",\\\"#8b008b\\\",\\\"#ff4500\\\",\\\"#8b4513\\\",\\\"#808000\\\",\\\"#483d8b\\\",\\\"#008000\\\",\\n\",\n \" \\\"#000080\\\",\\\"#9acd32\\\",\\\"#ffa500\\\",\\\"#ba55d3\\\",\\\"#00fa9a\\\",\\\"#dc143c\\\",\\\"#0000ff\\\",\\n\",\n \" \\\"#f08080\\\",\\\"#f0e68c\\\",\\\"#dda0dd\\\",\\\"#f2dcb3\\\",\\\"#f9cfcc\\\"]\\n\",\n \" return color_key[int(c)]\\n\",\n \"\\n\",\n \" def draw_layout(self,ax, doc, prob):\\n\",\n \" for spine in ax.spines.values():\\n\",\n \" spine.set_edgecolor('green')\\n\",\n \" spine.set_linewidth(1)\\n\",\n \" ax.set_xlim(0,self.n_col-1)\\n\",\n \" ax.set_ylim(0,self.n_row-1)\\n\",\n \" ax.get_xaxis().set_visible(False)\\n\",\n \" ax.get_yaxis().set_visible(False)\\n\",\n \" ax.invert_yaxis()\\n\",\n \"\\n\",\n \" for box in doc.T:\\n\",\n \" c = np.argmax(box[0:self.n_classes],axis=0)\\n\",\n \" x = np.argmax(box[self.n_classes:self.n_classes+self.n_col],axis=0)\\n\",\n \" y = np.argmax(box[self.n_classes+self.n_col:self.n_classes+self.n_col+self.n_row],axis=0)\\n\",\n \" w = np.argmax(box[self.n_classes+self.n_col+self.n_row:self.n_classes+2*self.n_col+self.n_row],axis=0)\\n\",\n \" h = np.argmax(box[self.n_classes+2*self.n_col+self.n_row:self.n_classes+2*self.n_col+2*self.n_row],axis=0)\\n\",\n \" r = plt.Rectangle((x,y),w,h, fc=self.get_color(c)+\\\"72\\\", ec=self.get_color(c),linewidth=1)\\n\",\n \" if c==self.n_classes-1:\\n\",\n \" break\\n\",\n \"\\n\",\n \" ax.add_patch(r)\\n\",\n \" \\n\",\n \"\\n\",\n \" def print_layouts(self,docs,min_boxes,beams_to_print,path=None,ratio_h_w=1.5):\\n\",\n \" plt.style.use('dark_background')\\n\",\n \" doc_num = docs[1]\\n\",\n \" docs = docs[0]\\n\",\n \" width = beams_to_print+3\\n\",\n \" height = ratio_h_w*(len(docs))+4\\n\",\n \" fig = plt.figure(figsize=(width,height),facecolor=\\\"#000000\\\",dpi=100)\\n\",\n \"\\n\",\n \" height_ratios = [0.8/(height),0.9/(height),(height-1.7)/height]\\n\",\n \" width_ratios = [(width-9)/(2*width),3/width,3/width,3/width,(width-9)/(2*width)]\\n\",\n \"\\n\",\n \" spec = gridspec.GridSpec(ncols=5, nrows=3,\\n\",\n \" width_ratios=width_ratios,\\n\",\n \" height_ratios=height_ratios,\\n\",\n \" wspace=5/width,left=0.05/width,right=(width-0.05)/width,top=0.98,bottom=0.02,hspace=0.05)\\n\",\n \" \\n\",\n \" ax = fig.add_subplot(spec[6])\\n\",\n \"\\n\",\n \" ax.plot(range(1,len(self.loss_his)+1),self.train_data_his,'-',color='red',linewidth=3)\\n\",\n \" ax.set_xlabel(\\\"Epochs\\\")\\n\",\n \" ax.set_title(\\\"Train Data\\\")\\n\",\n \" \\n\",\n \" ax = fig.add_subplot(spec[7])\\n\",\n \"\\n\",\n \" ax.plot(range(1,len(self.loss_his)+1),self.loss_his,'-',color='blue',linewidth=3)\\n\",\n \" ax.set_xlabel(\\\"Epochs\\\")\\n\",\n \" ax.set_title(\\\"KL Loss\\\")\\n\",\n \"\\n\",\n \" ax = fig.add_subplot(spec[8])\\n\",\n \"\\n\",\n \" ax.plot(range(1,len(self.loss_his)+1),self.lr_his,'-',color='green',linewidth=3)\\n\",\n \" ax.set_xlabel(\\\"Epochs\\\")\\n\",\n \" ax.set_title(\\\"LR\\\")\\n\",\n \" ax.set_yscale(\\\"log\\\")\\n\",\n \"\\n\",\n \" height_ratios = np.ones(len(docs)+2)*ratio_h_w/(height)\\n\",\n \" height_ratios[0] = 3.8/(height)\\n\",\n \" height_ratios[1] = 0.2/(height)\\n\",\n \"\\n\",\n \" spec = gridspec.GridSpec(ncols=1, nrows=len(docs)+2,\\n\",\n \" width_ratios=[1],\\n\",\n \" height_ratios=height_ratios,\\n\",\n \" wspace=0.05,left=0.02,right=0.98,top=0.98,bottom=0.02,hspace=0.05)\\n\",\n \"\\n\",\n \" ax = fig.add_subplot(spec[0])\\n\",\n \" ax.axis('off')\\n\",\n \" ax.invert_yaxis()\\n\",\n \" ax.text(0.5,0,\\\"Plot\\\",ha='center',va='bottom',fontsize=20)\\n\",\n \" doc_nums=\\\"\\\"\\n\",\n \" for i in doc_num:\\n\",\n \" doc_nums = doc_nums + \\\", \\\" + str(i)\\n\",\n \" doc_nums = doc_nums[2:]\\n\",\n \" doc_nums = \\\"Documents Predicted: \\\" + doc_nums\\n\",\n \" props = \\\"Classes: \\\" + str(self.n_classes-2) + \\\"; Epochs: \\\" + str(len(self.loss_his))\\n\",\n \" params = \\\"n_anchors = \\\" + str(self.n_anchors) + \\\"; d = \\\" + str(self.d) + \\\"; n_layers = \\\" + str(self.n_layers) + \\\"; n_heads = \\\" + str(self.n_heads) + \\\"; dff = \\\" + str(self.dff) + \\\"; dropout = \\\" + str(self.dropout)\\n\",\n \" ax.text(0.5,0.03,doc_nums+\\\"\\\\n\\\"+props+\\\"\\\\n\\\"+params,ha='center',va='top',fontsize=10)\\n\",\n \" ax.xaxis.set_visible(False)\\n\",\n \" ax.yaxis.set_visible(False)\\n\",\n \"\\n\",\n \" legend = []\\n\",\n \" legend.append(Patch(facecolor=self.get_color(0)+\\\"72\\\", label='',ec=self.get_color(0),linewidth=1))\\n\",\n \" for i in range(1,self.n_classes-1):\\n\",\n \" legend.append(Patch(facecolor=self.get_color(i)+\\\"72\\\", label=self.labels[i-1],ec=self.get_color(i),linewidth=1))\\n\",\n \" legend.append(Patch(facecolor=self.get_color(self.n_classes-1)+\\\"72\\\", label='',ec=self.get_color(self.n_classes-1),linewidth=1))\\n\",\n \" ax.legend(handles=legend,ncol=5,loc=8)\\n\",\n \"\\n\",\n \" height_ratios = np.ones(len(docs)+1)*ratio_h_w/(height)\\n\",\n \" height_ratios[0] = 4/(height)\\n\",\n \"\\n\",\n \" spec = gridspec.GridSpec(ncols=width, nrows=len(docs)+1,\\n\",\n \" width_ratios=np.ones(width),\\n\",\n \" height_ratios=height_ratios,\\n\",\n \" wspace=0.05,left=0.02,right=0.98,top=0.98,bottom=0.02,hspace=0.05)\\n\",\n \" \\n\",\n \" ax = fig.add_subplot(spec[0])\\n\",\n \" ax.axis('off')\\n\",\n \" ax.text(0.5,0,\\\"Ground Truth\\\",ha='center')\\n\",\n \"\\n\",\n \" ax = fig.add_subplot(spec[1])\\n\",\n \" ax.axis('off')\\n\",\n \" ax.text(0.5,0,\\\"Input\\\",ha='center')\\n\",\n \"\\n\",\n \" ax = fig.add_subplot(spec[2])\\n\",\n \" ax.axis('off')\\n\",\n \" ax.text(0.5,0,\\\"Most Probable\\\",ha='center')\\n\",\n \"\\n\",\n \" for i in range(1,beams_to_print+1):\\n\",\n \" ax = fig.add_subplot(spec[i+2])\\n\",\n \" ax.axis('off')\\n\",\n \" ax.text(0.5,0,\\\"Beam \\\"+str(i),ha='center')\\n\",\n \" \\n\",\n \" for input_count,doc_list in enumerate(docs):\\n\",\n \"\\n\",\n \" ax = fig.add_subplot(spec[(input_count+1)*width])\\n\",\n \" self.draw_layout(ax,doc_list[0][0],doc_list[0][1])\\n\",\n \" ax = fig.add_subplot(spec[(input_count+1)*width+1])\\n\",\n \" self.draw_layout(ax,doc_list[1][0],doc_list[1][1])\\n\",\n \" ax = fig.add_subplot(spec[(input_count+1)*width+2])\\n\",\n \" self.draw_layout(ax,doc_list[2][0],doc_list[2][1])\\n\",\n \"\\n\",\n \" doc_num=0\\n\",\n \" for doc in range(3,len(doc_list)):\\n\",\n \" if doc_num==beams_to_print:\\n\",\n \" break\\n\",\n \" if (len(doc_list[doc][0][0])>=min_boxes):\\n\",\n \" ax = fig.add_subplot(spec[(input_count+1)*width+doc])\\n\",\n \" self.draw_layout(ax,doc_list[doc][0],doc_list[doc][1])\\n\",\n \" doc_num = doc_num+1\\n\",\n \" if path!=None:\\n\",\n \" plt.savefig(path, facecolor=\\\"#000000\\\")\\n\",\n \" plt.show()\\n\",\n \" \\n\",\n \"\\n\",\n \" def layout_completion(self, initial_boxes_num=2, data_num_array=[0], beam_length=[1], max_boxes=10):\\n\",\n \" x = self.x_data[data_num_array,:,0:initial_boxes_num]\\n\",\n \" res = []\\n\",\n \"\\n\",\n \" for input_count,input in enumerate(x):\\n\",\n \" input = np.array([input])\\n\",\n \" most_prob_doc = [(input,1)]\\n\",\n \"\\n\",\n \" for step in range(max_boxes):\\n\",\n \" cur = most_prob_doc.pop(0)\\n\",\n \" pre = np.array([self.model(cur[0]).numpy()[0,:,-1]])\\n\",\n \" (box,p) = self.onehot(pre.T,[-1,-1,-1,-1,-1])\\n\",\n \" p=p*cur[1]\\n\",\n \" cur_box = np.array([box])\\n\",\n \" cur_doc = np.append(cur[0],cur_box,axis=2)\\n\",\n \"\\n\",\n \" most_prob_doc.append((cur_doc,p))\\n\",\n \"\\n\",\n \" if most_prob_doc[0][0][0,self.n_classes-1,-1]==1:\\n\",\n \" break\\n\",\n \"\\n\",\n \" docs = []\\n\",\n \" q = [(input,1)]\\n\",\n \" total_calc = 1\\n\",\n \" for step in range(max_boxes):\\n\",\n \" beam = 1\\n\",\n \" if step 1):\n cls_score = self.class_score(decoded)\n cls_prob = tf.sigmoid(tf.reshape(\n cls_score, [-1, self.n_component, self.n_class]))\n out = tf.concat([out, cls_prob], axis=-1)\n\n return out\n" }, { "path": "LayoutGAN/MNIST/mnist_train.py", "content": "import os\nimport time\nimport math\nfrom glob import glob\nimport tensorflow as tf\nimport numpy as np\nimport random\nfrom mnist_utils import *\nfrom mnist_modules import *\nimport matplotlib.pyplot as plt\nimport PIL\nfrom PIL import Image\nfrom PIL import ImageFont, ImageDraw\nfrom PIL import Image\n\nos.environ['CUDA_VISIBLE_DEVICES'] = '1'\n\n\nclass LayoutGAN(object):\n def __init__(self, geometric_dim=2, n_class=1, batch_size=64, n_component=128, layout_dim=(28, 28), d_lr=1e-5, g_lr=1.01e-5, update_ratio=2, clip_value=0.08568, dataset_name='default', dataset_path='./data/pre_data_cls.npy', checkpoint_dir=None, sample_dir=None):\n self.batch_size = batch_size\n self.n_component = n_component\n self.n_class = n_class\n self.geometric_dim = geometric_dim\n self.layout_dim = layout_dim\n self.dataset_name = dataset_name\n self.checkpoint_dir = checkpoint_dir\n self.data = np.load(dataset_path)\n self.build_model(d_lr, g_lr)\n self.sample_dir = sample_dir\n self.update_ratio = update_ratio\n self.clip_value = clip_value\n epoch_step = len(self.data) // self.batch_size\n dlr = tf.keras.optimizers.schedules.ExponentialDecay(\n 1e-5, epoch_step*20, 0.1, staircase=True, name=None)\n\n def build_model(self, dlr, g_lr):\n self.G = self.build_generator()\n self.D = self.build_discriminator()\n self.d_opt = tf.keras.optimizers.Adam(dlr)\n self.g_opt = tf.keras.optimizers.Adam(g_lr)\n\n def step(self, real_data, noise, training=True):\n with tf.GradientTape() as disc_tape:\n disc_loss = self.discriminator_loss(real_data, noise)\n\n if(training):\n gradients_of_discriminator = disc_tape.gradient(\n disc_loss, self.D.trainable_variables)\n self.d_opt.apply_gradients(\n zip(gradients_of_discriminator, self.D.trainable_variables))\n\n for i in range(2):\n with tf.GradientTape() as gen_tape:\n gen_loss = self.generator_loss(noise)\n\n if(training):\n gradients_of_generator = gen_tape.gradient(\n gen_loss, self.G.trainable_variables)\n self.g_opt.apply_gradients(\n zip(gradients_of_generator, self.G.trainable_variables))\n\n return gen_loss, disc_loss\n\n def train(self):\n epoch_step = len(self.data) // self.batch_size\n sample = self.data[0:self.batch_size]\n sample_inputs = np.array(sample).astype(np.float32)\n sample_inputs = sample_inputs * 28.0 / 27.0\n sample_z = np.random.normal(\n 0.5, 0.13, (self.batch_size, self.n_component, self.n_class, self.geometric_dim))\n counter = 1\n start_time = time.time()\n\n for epoch in range(150):\n np.random.shuffle(self.data)\n batch_idxs = len(self.data) // self.batch_size\n\n for idx in range(0, batch_idxs):\n batch = self.data[idx*self.batch_size:(idx+1)*self.batch_size]\n batch_images = np.array(batch).astype(np.float32)\n\n batch_images = batch_images * 28.0 / 27.0\n batch_z = np.random.normal(\n 0.5, 0.13, (self.batch_size, self.n_component, self.n_class, self.geometric_dim))\n g_loss, d_loss = self.step(batch_images, batch_z)\n counter += 1\n if np.mod(counter, 10) == 0:\n print(\"Epoch: [%2d] [%4d/%4d] time: %4.4f, d_loss: %.4f, g_loss: %.4f\"\n % (epoch, idx, batch_idxs, time.time()-start_time, d_loss, g_loss))\n\n if np.mod(counter, 1) == 0:\n samples = self.G(sample_z)\n g_loss, d_loss = self.step(\n sample_inputs, sample_z, training=False)\n samples = np.reshape(samples, (64, 128, 2))\n samples = 27.0 * samples\n\n img_all = np.zeros(\n (64, self.layout_dim[0], self.layout_dim[1], 3), dtype=np.uint8)\n rendered_layout = self.D.render(\n samples, self.layout_dim[0], self.layout_dim[1])\n img_list = []\n for img_ind in range(64):\n pointset = np.rint(\n samples[img_ind, :, :]).astype(np.int)\n pointset = pointset[~(pointset == 0).all(1)]\n\n img = np.zeros((28, 28), dtype=np.float32)\n img[pointset[:, 0], pointset[:, 1]] = 255\n img_list.append(img/255)\n img = Image.fromarray(img.astype('uint8'), 'L')\n img_all[img_ind, :, :, :] = np.array(\n img.convert('RGB'))\n img_all = np.squeeze(\n merge(img_all, image_manifold_size(samples.shape[0])))\n plt.imsave('{}/train_{:02d}_{:04d}.jpg'.format(self.sample_dir,\n epoch, idx), np.array(img_all, dtype=np.uint8))\n print(\"[Sample] d_loss: %.8f, g_loss: %.8f\" %\n (d_loss, g_loss))\n\n def render(self):\n pass\n\n def build_discriminator(self):\n return Discriminator(layout_dim=self.layout_dim, render=layout_point)\n\n def build_generator(self):\n return Generator(n_filters=512, output_dim=self.geometric_dim, n_component=self.n_component, n_class=self.n_class)\n\n def gradient_penalty(self, real, fake):\n alpha = tf.random.uniform(\n shape=[real.shape[0], 1, 1], minval=0.0, maxval=1.)\n interpolated = alpha * real + (1 - alpha) * fake\n with tf.GradientTape() as tape_p:\n tape_p.watch(interpolated)\n logit = self.D(interpolated)\n\n grad = tape_p.gradient(logit, interpolated)\n grad_norm = tf.norm(tf.reshape(grad, (real.shape[0], -1)), axis=1)\n\n return 10 * tf.reduce_mean(tf.square(grad_norm - 1.))\n\n def generator_loss(self, z):\n x = self.G(z, training=True)\n fake_score = self.D(x, training=True)\n g_loss = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(\n logits=fake_score, labels=tf.ones_like(tf.sigmoid(fake_score))))\n return g_loss\n\n def discriminator_loss(self, x, z):\n x_fake = self.G(z, training=True)\n true_score = self.D(x, training=True)\n fake_score = self.D(x_fake, training=True)\n d_loss_real = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(\n logits=true_score, labels=tf.ones_like(tf.sigmoid(true_score))))\n d_loss_fake = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(\n logits=fake_score, labels=tf.zeros_like(tf.sigmoid(fake_score))))\n d_loss = d_loss_real + d_loss_fake\n return d_loss\n\n\nbatch_size = 64\nn_component = 128\nn_class = 1\ngeometric_dim = 2\n# give approriate path\nsample_dir = \"../samples/MNIST_results\"\ngan = LayoutGAN(batch_size=batch_size, n_component=n_component, n_class=n_class, geometric_dim=geometric_dim,\n sample_dir=sample_dir, dataset_path=\"../data/mnist.npy\")\ngan.train()\n" }, { "path": "LayoutGAN/MNIST/mnist_utils.py", "content": "import numpy as np\nimport tensorflow as tf\nimport math\n\n\ndef merge(images, size):\n h, w = images.shape[1], images.shape[2]\n if (images.shape[3] in (3, 4)):\n c = images.shape[3] # size = 8 X 8 for 64 batch size\n img = np.zeros((h * size[0], w * size[1], c))\n for idx, image in enumerate(images):\n i = idx % size[1]\n j = idx // size[1]\n img[j * h:j * h + h, i * w:i * w + w, :] = image\n return img\n elif images.shape[3] == 1:\n img = np.zeros((h * size[0], w * size[1]))\n for idx, image in enumerate(images):\n i = idx % size[1]\n j = idx // size[1]\n img[j * h:j * h + h, i * w:i * w + w] = image[:, :, 0]\n return img\n\n\ndef image_manifold_size(num_images):\n manifold_h = int(np.floor(np.sqrt(num_images)))\n manifold_w = int(np.ceil(np.sqrt(num_images)))\n assert manifold_h * manifold_w == num_images\n return manifold_h, manifold_w\n\n\ndef layout_point(final_pred, output_height, output_width):\n bbox_pred = tf.reshape(final_pred, [64, 128, 2])\n\n x_r = tf.reshape(tf.range(output_width, dtype=tf.float32),\n [1, output_width, 1, 1])\n x_r = tf.reshape(tf.tile(x_r, [1, 1, output_width, 1]), [\n 1, output_width*output_width, 1, 1])\n x_r = tf.tile(x_r, [64, 1, 128, 1])\n\n y_r = tf.reshape(tf.range(output_height, dtype=tf.float32), [\n 1, 1, output_height, 1])\n y_r = tf.reshape(tf.tile(y_r, [1, output_height, 1, 1]), [\n 1, output_height*output_height, 1, 1])\n y_r = tf.tile(y_r, [64, 1, 128, 1])\n\n x_pred = tf.reshape(\n tf.slice(bbox_pred, [0, 0, 0], [-1, -1, 1]), [64, 1, 128, 1])\n x_pred = tf.tile(x_pred, [1, output_width*output_width, 1, 1])\n x_pred = (output_width-1.0) * x_pred\n\n y_pred = tf.reshape(\n tf.slice(bbox_pred, [0, 0, 1], [-1, -1, 1]), [64, 1, 128, 1])\n y_pred = tf.tile(y_pred, [1, output_height*output_height, 1, 1])\n y_pred = (output_height-1.0) * y_pred\n\n x_diff = tf.maximum(0.0, 1.0-tf.abs(x_r - x_pred))\n y_diff = tf.maximum(0.0, 1.0-tf.abs(y_r - y_pred))\n xy_diff = x_diff * y_diff\n\n xy_max = tf.nn.max_pool(xy_diff, ksize=[1, 1, 128, 1], strides=[\n 1, 1, 1, 1], padding='VALID')\n xy_max = tf.reshape(xy_max, [64, output_height, output_width, 1])\n\n return xy_max\n" }, { "path": "LayoutGAN/Publaynet/modules.py", "content": "import math\nfrom glob import glob\nimport tensorflow as tf\nfrom tensorflow.keras import initializers\nimport numpy as np\nimport random\nfrom utils import *\n\n\nclass RelationModule(tf.keras.Model):\n def __init__(self, channels=128, output_dim=128, key_dim=128, **kwargs):\n super(RelationModule, self).__init__(**kwargs)\n self.key_dim = channels\n self.output_dim = channels\n self.channels = channels\n self.key = tf.keras.layers.Conv2D(output_dim, (1, 1), strides=(1, 1), padding='same', kernel_initializer=initializers.TruncatedNormal(\n stddev=0.02, mean=0.0), bias_initializer=initializers.constant(0.0))\n self.query = tf.keras.layers.Conv2D(key_dim, (1, 1), strides=(1, 1), padding='same', kernel_initializer=initializers.TruncatedNormal(\n stddev=0.02, mean=0.0), bias_initializer=initializers.constant(0.0))\n self.value = tf.keras.layers.Conv2D(key_dim, (1, 1), strides=(1, 1), padding='same', kernel_initializer=initializers.TruncatedNormal(\n stddev=0.02, mean=0.0), bias_initializer=initializers.constant(0.0))\n self.projection = tf.keras.layers.Conv2D(channels, (1, 1), strides=(1, 1), padding='same', kernel_initializer=initializers.TruncatedNormal(\n stddev=0.02, mean=0.0), bias_initializer=initializers.constant(0.0))\n\n def call(self, inputs):\n f_k = tf.reshape(self.key(inputs), [\n inputs.shape[0], inputs.shape[1]*inputs.shape[2], self.key_dim])\n f_q = tf.reshape(self.query(inputs), [\n inputs.shape[0], inputs.shape[1]*inputs.shape[2], self.key_dim])\n f_q = tf.transpose(f_q, perm=[0, 2, 1])\n f_v = tf.reshape(self.value(inputs), [\n inputs.shape[0], inputs.shape[1]*inputs.shape[2], self.output_dim])\n\n attention_weight = tf.matmul(\n f_k, f_q)/(inputs.shape[1]*inputs.shape[2])\n out = tf.matmul(tf.transpose(attention_weight, perm=[0, 2, 1]), f_v)\n out = tf.reshape(\n out, [inputs.shape[0], inputs.shape[1], inputs.shape[2], self.output_dim])\n out = self.projection(out)\n return out\n\n\nclass Discriminator(tf.keras.Model):\n def __init__(self, n_filters=32, n_hidden=128, layout_dim=(28, 28), render=layout_bbox, **kwargs):\n super(Discriminator, self).__init__(**kwargs)\n self.layout_dim = layout_dim\n self.render = render\n self.act = tf.keras.layers.LeakyReLU(alpha=0.2)\n self.conv1 = tf.keras.layers.Conv2D(32, (5, 5), input_shape=layout_dim, strides=(\n 2, 2), padding='valid', kernel_initializer=initializers.TruncatedNormal(stddev=0.02, mean=0.0), bias_initializer=initializers.constant(0.0))\n self.bn1 = tf.keras.layers.BatchNormalization(\n epsilon=1e-5, momentum=0.9)\n\n self.conv2 = tf.keras.layers.Conv2D(32*2, (5, 5), strides=(2, 2), padding='valid', kernel_initializer=initializers.TruncatedNormal(\n stddev=0.02, mean=0.0), bias_initializer=initializers.constant(0.0))\n self.bn2 = tf.keras.layers.BatchNormalization(\n epsilon=1e-5, momentum=0.9)\n\n self.flatten = tf.keras.layers.Flatten()\n self.fc1 = tf.keras.layers.Dense(512, kernel_initializer=initializers.RandomNormal(\n stddev=0.02, mean=0.0), bias_initializer=initializers.constant(0.0))\n self.bn3 = tf.keras.layers.BatchNormalization(\n epsilon=1e-5, momentum=0.9)\n self.fc2 = tf.keras.layers.Dense(1, kernel_initializer=initializers.RandomNormal(\n stddev=0.02, mean=0.0), bias_initializer=initializers.constant(0.0))\n\n def call(self, inputs):\n x = self.render(inputs, self.layout_dim[0], self.layout_dim[1])\n x = self.act(self.bn1(self.conv1(x)))\n x = self.act(self.bn2(self.conv2(x)))\n x = self.flatten(x)\n x = self.act(self.bn3(self.fc1(x)))\n out = self.fc2(x)\n return out\n\n\nclass Generator(tf.keras.Model):\n def __init__(self, n_filters=128, output_dim=2, n_component=128, n_class=1, include_probability=False, **kwargs):\n super(Generator, self).__init__(**kwargs)\n self.n_filters = n_filters\n self.output_dim = output_dim\n self.n_component = n_component\n self.n_class = n_class\n self.include_probability = include_probability\n\n self.act = tf.keras.layers.ReLU()\n self.conv1_1 = tf.keras.layers.Conv2D(n_filters//4, (1, 1), input_shape=(self.n_component, 1, self.n_class+self.output_dim), strides=(\n 1, 1), padding='same', kernel_initializer=initializers.TruncatedNormal(stddev=0.02, mean=0.0), bias_initializer=initializers.constant(0.0))\n self.bn1_1 = tf.keras.layers.BatchNormalization(\n epsilon=1e-5, momentum=0.9)\n self.conv1_2 = tf.keras.layers.Conv2D(n_filters//16, (1, 1), strides=(1, 1), padding='same',\n kernel_initializer=initializers.TruncatedNormal(stddev=0.02, mean=0.0), bias_initializer=initializers.constant(0.0))\n self.bn1_2 = tf.keras.layers.BatchNormalization(\n epsilon=1e-5, momentum=0.9)\n self.conv1_3 = tf.keras.layers.Conv2D(n_filters//16, (1, 1), strides=(1, 1), padding='same',\n kernel_initializer=initializers.TruncatedNormal(stddev=0.02, mean=0.0), bias_initializer=initializers.constant(0.0))\n self.bn1_3 = tf.keras.layers.BatchNormalization(\n epsilon=1e-5, momentum=0.9)\n self.conv1_4 = tf.keras.layers.Conv2D(n_filters//4, (1, 1), strides=(1, 1), padding='same',\n kernel_initializer=initializers.TruncatedNormal(stddev=0.02, mean=0.0), bias_initializer=initializers.constant(0.0))\n self.bn1_4 = tf.keras.layers.BatchNormalization(\n epsilon=1e-5, momentum=0.9)\n\n self.relation1 = RelationModule(\n channels=n_filters//4, output_dim=n_filters//4, key_dim=n_filters//4)\n self.relation2 = RelationModule(\n channels=n_filters//4, output_dim=n_filters//4, key_dim=n_filters//4)\n self.relation3 = RelationModule(\n channels=n_filters, output_dim=n_filters, key_dim=n_filters)\n self.relation4 = RelationModule(\n channels=n_filters, output_dim=n_filters, key_dim=n_filters)\n\n self.bn_x1 = tf.keras.layers.BatchNormalization(\n epsilon=1e-5, momentum=0.9)\n self.bn_x2 = tf.keras.layers.BatchNormalization(\n epsilon=1e-5, momentum=0.9)\n self.bn_x3 = tf.keras.layers.BatchNormalization(\n epsilon=1e-5, momentum=0.9)\n self.bn_x4 = tf.keras.layers.BatchNormalization(\n epsilon=1e-5, momentum=0.9)\n\n self.bn_x5 = tf.keras.layers.BatchNormalization(\n epsilon=1e-5, momentum=0.9)\n self.bn_x6 = tf.keras.layers.BatchNormalization(\n epsilon=1e-5, momentum=0.9)\n self.bn_x7 = tf.keras.layers.BatchNormalization(\n epsilon=1e-5, momentum=0.9)\n self.bn_x8 = tf.keras.layers.BatchNormalization(\n epsilon=1e-5, momentum=0.9)\n\n self.conv2_1 = tf.keras.layers.Conv2D(n_filters, (1, 1), strides=(1, 1), padding='same', kernel_initializer=initializers.TruncatedNormal(\n stddev=0.02, mean=0.0), bias_initializer=initializers.constant(0.0))\n self.bn2_1 = tf.keras.layers.BatchNormalization(\n epsilon=1e-5, momentum=0.9)\n self.conv2_2 = tf.keras.layers.Conv2D(n_filters//4, (1, 1), strides=(1, 1), padding='same',\n kernel_initializer=initializers.TruncatedNormal(stddev=0.02, mean=0.0), bias_initializer=initializers.constant(0.0))\n self.bn2_2 = tf.keras.layers.BatchNormalization(\n epsilon=1e-5, momentum=0.9)\n self.conv2_3 = tf.keras.layers.Conv2D(n_filters//4, (1, 1), strides=(1, 1), padding='same',\n kernel_initializer=initializers.TruncatedNormal(stddev=0.02, mean=0.0), bias_initializer=initializers.constant(0.0))\n self.bn2_3 = tf.keras.layers.BatchNormalization(\n epsilon=1e-5, momentum=0.9)\n self.conv2_4 = tf.keras.layers.Conv2D(n_filters, (1, 1), strides=(1, 1), padding='same', kernel_initializer=initializers.TruncatedNormal(\n stddev=0.02, mean=0.0), bias_initializer=initializers.constant(0.0))\n self.bn2_4 = tf.keras.layers.BatchNormalization(\n epsilon=1e-5, momentum=0.9)\n self.geometric_param = tf.keras.layers.Conv2D(output_dim, (1, 1), strides=(\n 1, 1), padding='same', kernel_initializer=initializers.TruncatedNormal(stddev=0.001, mean=0.0), bias_initializer=initializers.constant(0.0))\n self.class_score = tf.keras.layers.Conv2D(n_class, (1, 1), strides=(1, 1), padding='same', kernel_initializer=initializers.TruncatedNormal(\n stddev=0.02, mean=0.0), bias_initializer=initializers.constant(0.0))\n\n def call(self, x):\n x = tf.reshape(x, [x.shape[0], self.n_component,\n 1, self.n_class+self.output_dim])\n h1_0 = self.bn1_1(self.conv1_1(x))\n h1_1 = self.act(self.bn1_2(self.conv1_2(x)))\n h1_2 = self.act(self.bn1_3(self.conv1_3(h1_1)))\n h1_3 = self.bn1_4(self.conv1_4(h1_2))\n embedding = self.act(tf.add(h1_0, h1_3))\n embedding = tf.reshape(\n embedding, [x.shape[0], self.n_component, 1, 256])\n\n context = self.act(self.bn_x2(\n tf.add(embedding, self.bn_x1(self.relation1(embedding)))))\n context = self.act(self.bn_x4(\n tf.add(context, self.bn_x3(self.relation2(context)))))\n\n h2_0 = self.bn2_1(self.conv2_1(context))\n h2_1 = self.act(self.bn2_2(self.conv2_2(h2_0)))\n h2_2 = self.act(self.bn2_3(self.conv2_3(h2_1)))\n h2_3 = self.bn2_4(self.conv2_4(h2_2))\n decoded = self.act(tf.add(h2_0, h2_3))\n\n decoded = self.act(self.bn_x6(\n tf.add(decoded, self.bn_x5(self.relation3(decoded)))))\n decoded = self.act(self.bn_x8(\n tf.add(decoded, self.bn_x7(self.relation4(decoded)))))\n\n out = self.geometric_param(decoded)\n out = tf.sigmoid(tf.reshape(\n out, [-1, self.n_component, self.output_dim]))\n\n cls_score = self.class_score(decoded)\n cls_prob = tf.sigmoid(tf.reshape(\n cls_score, [-1, self.n_component, self.n_class]))\n final_pred = tf.concat([out, cls_prob], axis=-1)\n return final_pred\n" }, { "path": "LayoutGAN/Publaynet/train.py", "content": "import os\nimport time\nimport math\nfrom glob import glob\nimport tensorflow as tf\nimport numpy as np\nimport random\nfrom utils import *\nfrom modules import *\nimport matplotlib.pyplot as plt\nfrom tensorflow.keras import initializers\n\n\nos.environ['CUDA_VISIBLE_DEVICES'] = '1'\n\n\nclass LayoutGAN(object):\n def __init__(self, geometric_dim=2, n_class=1, batch_size=64, n_component=128, layout_dim=(28, 28), d_lr=1e-5, g_lr=1e-5, update_ratio=2, clip_value=0.1, dataset_name='default', dataset_path='./data/pre_data_cls.npy', checkpoint_dir=None, sample_dir=None):\n self.batch_size = batch_size\n self.n_component = n_component\n self.n_class = n_class\n self.geometric_dim = geometric_dim\n self.layout_dim = layout_dim\n self.dataset_name = dataset_name\n self.checkpoint_dir = checkpoint_dir\n self.data = np.load(dataset_path)\n self.data = self.data[:70000]\n self.build_model(d_lr, g_lr)\n self.sample_dir = sample_dir\n self.update_ratio = update_ratio\n self.clip_value = clip_value\n self.epochs = 50\n\n def build_model(self, d_lr, g_lr):\n self.G = self.build_generator()\n self.D = self.build_discriminator()\n epoch_step = len(self.data) // self.batch_size\n dlr = tf.keras.optimizers.schedules.ExponentialDecay(\n d_lr, decay_steps=20*epoch_step, decay_rate=0.1, staircase=True)\n self.d_opt = tf.keras.optimizers.Adam(dlr)\n self.g_opt = tf.keras.optimizers.Adam(dlr)\n\n def step(self, real_data, noise, training=True, step=0):\n with tf.GradientTape() as disc_tape:\n disc_loss = self.discriminator_loss(real_data, noise)\n if(training):\n gradients_of_discriminator = disc_tape.gradient(\n disc_loss, self.D.trainable_variables)\n self.d_opt.apply_gradients(\n zip(gradients_of_discriminator, self.D.trainable_variables))\n\n for i in range(self.update_ratio):\n with tf.GradientTape() as gen_tape:\n gen_loss = self.generator_loss(noise)\n if(training):\n gradients_of_generator = gen_tape.gradient(\n gen_loss, self.G.trainable_variables)\n self.g_opt.apply_gradients(\n zip(gradients_of_generator, self.G.trainable_variables))\n\n return gen_loss, disc_loss\n\n def train(self):\n epoch_step = len(self.data) // self.batch_size\n sample = self.data[0:self.batch_size]\n sample_inputs = np.array(sample).astype(np.float32)\n sample_z_bbox = np.random.normal(0.5, 0.15, (self.batch_size, 9, 4))\n sample_z_cls = np.identity(\n 5)[np.random.randint(5, size=(self.batch_size, 9))]\n sample_z = np.concatenate([sample_z_bbox, sample_z_cls], axis=-1)\n counter = 1\n start_time = time.time()\n\n for epoch in range(self.epochs):\n np.random.shuffle(self.data)\n batch_idxs = len(self.data) // self.batch_size\n\n for idx in range(0, batch_idxs):\n batch = self.data[idx*self.batch_size:(idx+1)*self.batch_size]\n batch_images = np.array(batch).astype(np.float32)\n\n batch_z_bbox = np.random.normal(\n 0.5, 0.15, (self.batch_size, 9, 4))\n batch_z_cls = np.identity(\n 5)[np.random.randint(5, size=(self.batch_size, 9))]\n batch_z = np.concatenate([batch_z_bbox, batch_z_cls], axis=-1)\n\n g_loss, d_loss = self.step(batch_images, batch_z, step=idx)\n counter += 1\n if np.mod(counter, 50) == 0:\n\n current_decayed_lr = self.d_opt._decayed_lr(\n tf.float32).numpy()\n print(\"Epoch: [%2d] [%4d/%4d] time: %4.4f, lr:%.3E, d_loss: %.4f, g_loss: %.4f\"\n % (epoch, idx, batch_idxs, time.time()-start_time, current_decayed_lr, d_loss, g_loss))\n\n if np.mod(counter, 500) == 0:\n G_samples = self.G(sample_z, training=False)\n path = '{}/train_{:02d}_{:04d}_{:2.4f}_{:2.4f}.jpg'.format(\n self.sample_dir, epoch, idx, d_loss, g_loss)\n change = convert_to_cxywh(np.array(G_samples))\n plot_layouts(change, colors=colors,\n class_names=class_names, path=path)\n g_loss, d_loss = self.step(\n sample_inputs, sample_z, training=False)\n print(\"[Sample] d_loss: %.8f, g_loss: %.8f\" %\n (d_loss, g_loss))\n\n def render(self):\n pass\n\n def build_discriminator(self):\n return Discriminator(layout_dim=self.layout_dim, render=layout_bbox)\n\n def build_generator(self):\n return Generator(n_filters=1024, output_dim=self.geometric_dim, n_component=self.n_component, n_class=self.n_class)\n\n def generator_loss(self, z):\n x = self.G(z, training=True)\n fake_score = self.D(x, training=True)\n g_loss = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(\n logits=fake_score, labels=tf.ones_like(tf.sigmoid(fake_score))))\n return g_loss\n\n def discriminator_loss(self, x, z):\n x_fake = self.G(z, training=True)\n true_score = self.D(x, training=True)\n fake_score = self.D(x_fake, training=True)\n d_loss_real = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(\n logits=true_score, labels=tf.ones_like(tf.sigmoid(true_score))))\n d_loss_fake = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(\n logits=fake_score, labels=tf.zeros_like(tf.sigmoid(fake_score))))\n d_loss = d_loss_real + d_loss_fake\n return d_loss\n\n\nif __name__ == '__main__':\n batch_size = 64\n n_component = 9\n n_class = 5\n geometric_dim = 4\n gan = LayoutGAN(batch_size=batch_size, n_component=n_component,\n n_class=n_class, layout_dim=(60, 40),\n geometric_dim=geometric_dim,\n sample_dir=\"./samples/publaynet_results\",\n dataset_path=\"./data/sorted_c1publay.npy\")\n\n gan.train()\n" }, { "path": "LayoutGAN/Publaynet/utils.py", "content": "import numpy as np\nimport tensorflow as tf\nimport matplotlib.pyplot as plt\nimport matplotlib.patches as patches\nfrom matplotlib.patches import Patch\nplt.style.use('dark_background')\n\n\ndef convert_to_cxywh(data):\n bboxes = data[..., 0:4]\n labels = data[..., 4:]\n mask = np.zeros_like(data[..., 3:4])\n labels = np.concatenate((mask, labels), axis=2)\n labels = np.argmax(labels, axis=2)\n class_info = np.expand_dims(labels, axis=2)\n cxywh = np.concatenate((class_info, bboxes), axis=2)\n cxywh[..., 1] = cxywh[..., 1] - cxywh[..., 3]/2\n cxywh[..., 2] = cxywh[..., 2] - cxywh[..., 4]/2\n return cxywh\n\n\ndef generate_colors(class_names=None, n_class=50):\n cmap = [\"\", \"#dc143c\", \"#ffff00\", \"#00ff00\", \"#ff00ff\", \"#1e90ff\", \"#fff5ee\",\n \"#00ffff\", \"#8b008b\", \"#ff4500\", \"#8b4513\", \"#808000\", \"#483d8b\",\n \"#008000\", \"#000080\", \"#9acd32\", \"#ffa500\", \"#ba55d3\", \"#00fa9a\",\n \"#dc143c\", \"#0000ff\", \"#f08080\", \"#f0e68c\", \"#dda0dd\", \"#ff1493\"]\n colors = dict()\n if class_names == None:\n class_names = []\n for i in range(n_class):\n class_names.append('class'+str(i+1))\n for i in range(n_class):\n colors[class_names[i]] = cmap[i]\n return colors\n\n\nclass_names = ['None', 'Text', 'Title', 'List', 'Table', 'Figure']\ncolors = generate_colors(n_class=6, class_names=class_names)\n\n\ndef plot_layouts(pred, colors, class_names, path=\"\"):\n height = 15\n width = 9\n fig = plt.figure(figsize=(width, height), dpi=50, facecolor=(0, 0, 0))\n height_ratio = [0.25, 1, 1, 1, 1, 1]\n grid = plt.GridSpec(6, 4,\n hspace=0.05, wspace=0.05,\n height_ratios=height_ratio,\n left=0.02, right=0.98, top=0.98, bottom=0.02)\n index = 0\n legend = []\n ax = fig.add_subplot(grid[index: index+4])\n index += 4\n for i in range(1, 6):\n legend.append(Patch(facecolor=colors[class_names[i]]+\"40\",\n edgecolor=colors[class_names[i]],\n label=class_names[i]))\n\n ax.legend(handles=legend, ncol=3, loc=8, fontsize=25, facecolor=(0, 0, 0))\n ax.axis('off')\n\n for i in range(16):\n ax = fig.add_subplot(grid[index])\n index += 1\n\n data = pred[i]\n rect1 = patches.Rectangle((0, 0), 180, 240)\n rect1.set_color((0, 0, 0, 1))\n ax.add_patch(rect1)\n for box in data:\n\n c, x, y, w, h = box\n if c == 0:\n continue\n x = x*180\n y = y*240\n w = w*180\n h = h*240\n rect = patches.Rectangle((x, y), w, h, linewidth=2)\n rect.set_color(colors[class_names[int(c)]]+\"00\")\n rect.set_linestyle('-')\n rect.set_edgecolor(colors[class_names[int(c)]])\n ax.add_patch(rect)\n ax.plot()\n ax.set_facecolor((0, 0, 0))\n for spine in ax.spines.values():\n spine.set_edgecolor('green')\n spine.set_linewidth(2)\n ax.invert_yaxis()\n ax.set_xticks([])\n ax.set_yticks([])\n plt.savefig(path, facecolor=(0, 0, 0))\n\n\ndef layout_bbox(final_pred, output_height, output_width):\n final_pred = tf.reshape(final_pred, [64, 9, 9])\n bbox_reg = tf.slice(final_pred, [0, 0, 0], [-1, -1, 4])\n cls_prob = tf.slice(final_pred, [0, 0, 4], [-1, -1, 5])\n\n bbox_reg = tf.reshape(bbox_reg, [64, 9, 4])\n\n x_c = tf.slice(bbox_reg, [0, 0, 0], [-1, -1, 1]) * output_width\n y_c = tf.slice(bbox_reg, [0, 0, 1], [-1, -1, 1]) * output_height\n w = tf.slice(bbox_reg, [0, 0, 2], [-1, -1, 1]) * output_width\n h = tf.slice(bbox_reg, [0, 0, 3], [-1, -1, 1]) * output_height\n\n x1 = x_c - 0.5*w\n x2 = x_c + 0.5*w\n y1 = y_c - 0.5*h\n y2 = y_c + 0.5*h\n\n xt = tf.reshape(tf.range(output_width, dtype=tf.float32), [1, 1, 1, -1])\n xt = tf.reshape(tf.tile(xt, [64, 9, output_height, 1]), [64, 9, -1])\n\n yt = tf.reshape(tf.range(output_height, dtype=tf.float32), [1, 1, -1, 1])\n yt = tf.reshape(tf.tile(yt, [64, 9, 1, output_width]), [64, 9, -1])\n\n x1_diff = tf.reshape(xt-x1, [64, 9, output_height, output_width, 1])\n y1_diff = tf.reshape(yt-y1, [64, 9, output_height, output_width, 1])\n x2_diff = tf.reshape(x2-xt, [64, 9, output_height, output_width, 1])\n y2_diff = tf.reshape(y2-yt, [64, 9, output_height, output_width, 1])\n\n x1_line = tf.nn.relu(1.0 - tf.abs(x1_diff)) * tf.minimum(\n tf.nn.relu(y1_diff), 1.0) * tf.minimum(tf.nn.relu(y2_diff), 1.0)\n x2_line = tf.nn.relu(1.0 - tf.abs(x2_diff)) * tf.minimum(\n tf.nn.relu(y1_diff), 1.0) * tf.minimum(tf.nn.relu(y2_diff), 1.0)\n y1_line = tf.nn.relu(1.0 - tf.abs(y1_diff)) * tf.minimum(\n tf.nn.relu(x1_diff), 1.0) * tf.minimum(tf.nn.relu(x2_diff), 1.0)\n y2_line = tf.nn.relu(1.0 - tf.abs(y2_diff)) * tf.minimum(\n tf.nn.relu(x1_diff), 1.0) * tf.minimum(tf.nn.relu(x2_diff), 1.0)\n\n xy_max = tf.reduce_max(tf.concat(\n [x1_line, x2_line, y1_line, y2_line], axis=-1), axis=-1, keepdims=True)\n\n spatial_prob = tf.multiply(\n tf.tile(xy_max, [1, 1, 1, 1, 5]), tf.reshape(cls_prob, [64, 9, 1, 1, 5]))\n spatial_prob_max = tf.reduce_max(spatial_prob, axis=1, keepdims=False)\n\n return spatial_prob_max\n" }, { "path": "LayoutGAN/README.md", "content": "# LayoutGAN \nThis repository provides implementation of \"LayoutGAN: Generating Graphic Layouts with Wireframe Discriminators\"\nhttps://arxiv.org/abs/1901.06767 in Tensorflow 2.0. \n\n## Getting Started\n\n\"Open\n\nClick on the above Badge to quickstart the LayoutGAN in google colab\n\n### Architecture \n![](demo/layoutgan.png)\n\n## Prerequisites\n- Python 3.8\n- Tensorflow 2.5.0\n\n## MNIST Generation\n1. Use [MNIST Dataset](https://drive.google.com/file/d/1qtBnEWsaKXeynOCUFHB7H9eqU6bDLJAV/view?usp=sharing). Download and put it in the data folder.\n2. Dataset has dimension of (70,000, 9, 9), which consists of 70,000 samples, which contains 128 foreground pixels cordinates (N, 128, X, Y).\n3. Run `python MNIST/mnist_train.py` to train a model.\n4. Predictions will be saved in `samples/MNIST_results`.\n### Results on MNIST\n![](demo/mnist_obtained.jpeg)\n\n## Document Layout Generation\n1. Use [PubLayNet Dataset](https://drive.google.com/file/d/1YQKyASvGDNUTJnE1x-Q2ZhhiY0VFj7oZ/view?usp=sharing). Download and put it in data folder.\n2. Dataset has dimension of (72499, 9, 9), which consists of 72,499 single column layout samples extracted from PubLayNet dataset, which contains atmost 9 bounding boxes,for each bounding box first four elements are dimensions of bounding boxes (X_centre, Y_centre, Width, Height).\n3. Run `python publaynet/train.py` to train a model.\n4. Predictions will be saved in `samples/publaynet_results`.\n### Results on single column layouts\n\n\n## Related repositories\nSome codes are implemented from\nhttps://github.com/JiananLi2016/LayoutGAN-Tensorflow\n\n" }, { "path": "LayoutGAN/data/.gitkeep", "content": "" }, { "path": "LayoutGAN/demo/.gitkeep", "content": "" }, { "path": "LayoutGAN/samples/MNIST_results/.gitkeep", "content": "" }, { "path": "LayoutGAN/samples/publaynet_results/.gitkeep", "content": "" }, { "path": "LayoutVAE/Notebook/LayoutVAE_Final.ipynb", "content": "{\n \"nbformat\": 4,\n \"nbformat_minor\": 0,\n \"metadata\": {\n \"kernelspec\": {\n \"language\": \"python\",\n \"display_name\": \"Python 3\",\n \"name\": \"python3\"\n },\n \"language_info\": {\n \"name\": \"python\",\n \"version\": \"3.7.10\",\n \"mimetype\": \"text/x-python\",\n \"codemirror_mode\": {\n \"name\": \"ipython\",\n \"version\": 3\n },\n \"pygments_lexer\": \"ipython3\",\n \"nbconvert_exporter\": \"python\",\n \"file_extension\": \".py\"\n },\n \"colab\": {\n \"name\": \"LayoutVAE-Final.ipynb\",\n \"provenance\": [],\n \"collapsed_sections\": []\n },\n \"accelerator\": \"GPU\"\n },\n \"cells\": [\n {\n \"cell_type\": \"code\",\n \"metadata\": {\n \"id\": \"aoSIFtCBTLAE\"\n },\n \"source\": [\n \"from google.colab import drive\\n\",\n \"drive.mount('/content/drive')\"\n ],\n \"execution_count\": null,\n \"outputs\": []\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {\n \"id\": \"b3hWjvVkcw7G\"\n },\n \"source\": [\n \"# Imports\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"metadata\": {\n \"id\": \"Du4WO_dO0W2o\",\n \"execution\": {\n \"iopub.status.busy\": \"2021-07-22T09:46:51.330856Z\",\n \"iopub.execute_input\": \"2021-07-22T09:46:51.331216Z\",\n \"iopub.status.idle\": \"2021-07-22T09:46:51.719259Z\",\n \"shell.execute_reply.started\": \"2021-07-22T09:46:51.331182Z\",\n \"shell.execute_reply\": \"2021-07-22T09:46:51.718319Z\"\n },\n \"trusted\": true\n },\n \"source\": [\n \"from __future__ import division\\n\",\n \"import torch as T\\n\",\n \"import torch.functional as F\\n\",\n \"import math\\n\",\n \"import PIL\\n\",\n \"import numpy as np\\n\",\n \"import pandas as pd\\n\",\n \"import matplotlib.pyplot as plt\\n\",\n \"import matplotlib.patches as patches\\n\",\n \"import matplotlib.gridspec as gridspec\\n\",\n \"from matplotlib.patches import Patch\\n\",\n \"from torch.nn import Sequential , Linear , ReLU , PoissonNLLLoss, LSTM\\n\",\n \"from torch.autograd import Variable\\n\",\n \"from torch.distributions import Normal, MultivariateNormal, Poisson, kl_divergence \\n\",\n \"from PIL import Image,ImageFont, ImageDraw\\n\",\n \"plt.style.use('dark_background')\\n\",\n \"T.set_default_tensor_type('torch.cuda.FloatTensor')\"\n ],\n \"execution_count\": null,\n \"outputs\": []\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {\n \"id\": \"3599a3ADc2aR\"\n },\n \"source\": [\n \"# Paths\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"metadata\": {\n \"id\": \"2DZL1CwwetFu\"\n },\n \"source\": [\n \"root = \\\"\\\"\\n\",\n \"DATA_PATH = root + \\\"\\\"\\n\",\n \"SAVE_MODEL_PATH = root + \\\"\\\"\\n\",\n \"SAVE_LOG_PATH = root + \\\"\\\"\\n\",\n \"SAVE_OUTPUT_PATH = root + \\\"\\\"\\n\",\n \"CVAE_PATH = root + \\\"\\\"\"\n ],\n \"execution_count\": null,\n \"outputs\": []\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {\n \"id\": \"c0mBZWyzc4S6\"\n },\n \"source\": [\n \"# Model Architectures\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"metadata\": {\n \"id\": \"89YDE-Eg9VBh\",\n \"execution\": {\n \"iopub.status.busy\": \"2021-07-22T12:06:02.658911Z\",\n \"iopub.execute_input\": \"2021-07-22T12:06:02.659249Z\",\n \"iopub.status.idle\": \"2021-07-22T12:06:02.678743Z\",\n \"shell.execute_reply.started\": \"2021-07-22T12:06:02.659216Z\",\n \"shell.execute_reply\": \"2021-07-22T12:06:02.677835Z\"\n },\n \"trusted\": true\n },\n \"source\": [\n \"class fcblock(T.nn.Module):\\n\",\n \" def __init__(self, n_class):\\n\",\n \" super(fcblock, self).__init__()\\n\",\n \" self.seq = Sequential(\\n\",\n \" Linear(n_class,128),\\n\",\n \" ReLU(),\\n\",\n \" Linear(128,128),\\n\",\n \" ReLU(),\\n\",\n \" )\\n\",\n \" def forward(self,inputs):\\n\",\n \" out = self.seq(inputs)\\n\",\n \" return out\\n\",\n \"\\n\",\n \"class Embeder(T.nn.Module):\\n\",\n \" def __init__(self,n_class):\\n\",\n \" super(Embeder,self).__init__()\\n\",\n \" \\n\",\n \" self.fcb1 = fcblock(n_class)\\n\",\n \" self.fcb2 = fcblock(n_class)\\n\",\n \" self.fcb3 = fcblock(n_class)\\n\",\n \" self.fc = Linear(128*3,128)\\n\",\n \"\\n\",\n \" \\n\",\n \" def forward(self,inputs):\\n\",\n \" in1,in2,in3 = inputs\\n\",\n \" in1 = self.fcb1(in1)\\n\",\n \" in2 = self.fcb2(in2)\\n\",\n \" in3 = self.fcb3(in3)\\n\",\n \" out = T.cat((in1,in2,in3),1)\\n\",\n \" out = self.fc(out)\\n\",\n \" return out\\n\",\n \"\\n\",\n \"class Encoder(T.nn.Module):\\n\",\n \" def __init__(self, in_dim=1 ,latent_dim=32):\\n\",\n \" super(Encoder,self).__init__()\\n\",\n \" self.act = ReLU()\\n\",\n \" self.fc1 = Linear(in_dim,128)\\n\",\n \" self.fc2 = Linear(128,128)\\n\",\n \" self.fc3 = Linear(256,latent_dim)\\n\",\n \" self.fc4 = Linear(latent_dim,latent_dim)\\n\",\n \" self.fc5 = Linear(latent_dim,latent_dim)\\n\",\n \" \\n\",\n \" def forward(self,inputs):\\n\",\n \" in1,in2 = inputs\\n\",\n \" out = self.fc1(in1)\\n\",\n \" out = self.act(out)\\n\",\n \" out = self.fc2(out)\\n\",\n \" out = T.cat((out,in2),1)\\n\",\n \" out = self.fc3(out)\\n\",\n \" out = self.act(out)\\n\",\n \" mu = self.fc4(out)\\n\",\n \" logvar = self.fc5(out)\\n\",\n \" return mu,logvar\\n\",\n \"\\n\",\n \"\\n\",\n \"class Prior(T.nn.Module):\\n\",\n \" def __init__(self,latent_dim=32):\\n\",\n \" super(Prior,self).__init__()\\n\",\n \" \\n\",\n \" self.act = ReLU()\\n\",\n \" self.fc1 = Linear(128,latent_dim)\\n\",\n \" self.fc2 = Linear(latent_dim,latent_dim)\\n\",\n \" self.fc3 = Linear(latent_dim,latent_dim)\\n\",\n \" \\n\",\n \" def forward(self,inputs):\\n\",\n \" out = inputs\\n\",\n \" out = self.fc1(out)\\n\",\n \" out = self.act(out)\\n\",\n \" mu = self.fc2(out)\\n\",\n \" logvar = self.fc3(out) \\n\",\n \" return mu,logvar\\n\",\n \"\\n\",\n \"class Decoder(T.nn.Module):\\n\",\n \" def __init__(self,output_dim,latent_dim=32):\\n\",\n \" super(Decoder,self).__init__()\\n\",\n \" self.act = ReLU()\\n\",\n \" self.fc1 = Linear(128+latent_dim,128)\\n\",\n \" self.fc2 = Linear(128,64)\\n\",\n \" self.fc3 = Linear(64,output_dim)\\n\",\n \" \\n\",\n \" def forward(self,inputs):\\n\",\n \" in1,in2 = inputs\\n\",\n \" out = T.cat((in1,in2),1)\\n\",\n \" out = self.fc1(out)\\n\",\n \" out = self.act(out)\\n\",\n \" out = self.fc2(out)\\n\",\n \" out = self.act(out)\\n\",\n \" out = self.fc3(out)\\n\",\n \" return out\"\n ],\n \"execution_count\": null,\n \"outputs\": []\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {\n \"id\": \"lqWGZGytlSi7\"\n },\n \"source\": [\n \"# Loss Function for Countvae\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"metadata\": {\n \"id\": \"zwTGSyrelNrV\",\n \"execution\": {\n \"iopub.status.busy\": \"2021-07-22T12:08:51.811748Z\",\n \"iopub.execute_input\": \"2021-07-22T12:08:51.812140Z\",\n \"iopub.status.idle\": \"2021-07-22T12:08:51.822374Z\",\n \"shell.execute_reply.started\": \"2021-07-22T12:08:51.812099Z\",\n \"shell.execute_reply\": \"2021-07-22T12:08:51.821128Z\"\n },\n \"trusted\": true\n },\n \"source\": [\n \"class ELBOLoss(T.nn.Module):\\n\",\n \"\\n\",\n \" def __init__(self):\\n\",\n \" super(ELBOLoss,self).__init__()\\n\",\n \" \\n\",\n \" def forward(self,inputs):\\n\",\n \" mu1, logvar1, mu2, logvar2 , in1, in2 = inputs\\n\",\n \"\\n\",\n \" mask = (in2>0)+0.0\\n\",\n \" in2 = in2-mask\\n\",\n \"\\n\",\n \" '''KL Divergence'''\\n\",\n \" kl = 0.5 * T.sum((logvar2 - logvar1) - 1 + (logvar1.exp() + (mu2 - mu1).pow(2) )/logvar2.exp() , dim = 1).mean()\\n\",\n \" \\n\",\n \" '''Poisson Negative Log Likelihood'''\\n\",\n \" pnll = PoissonNLLLoss()(in1,in2)\\n\",\n \"\\n\",\n \" loss = kl+pnll\\n\",\n \" \\n\",\n \" return loss, pnll , kl\\n\",\n \" \\n\",\n \" \\n\",\n \"class Reparamatrize(T.nn.Module):\\n\",\n \" \\n\",\n \" def __init__(self):\\n\",\n \" super(Reparamatrize,self).__init__()\\n\",\n \" \\n\",\n \" def forward(self,inputs):\\n\",\n \" \\n\",\n \" mu , logvar = inputs\\n\",\n \" '''\\n\",\n \" mu = mean \\n\",\n \" logvar = log of diagonal elements of covariance matrix\\n\",\n \" '''\\n\",\n \" # Covarince Matrix\\n\",\n \" covar = T.diag_embed(T.exp(logvar/2), dim1=-2,dim2=-1)\\n\",\n \"\\n\",\n \" # Multivariate Normal Distribution\\n\",\n \" p = MultivariateNormal(mu,covar)\\n\",\n \" z_latent = p.rsample().float()\\n\",\n \" return z_latent\\n\",\n \"\\n\",\n \"class Sampling(T.nn.Module):\\n\",\n \"\\n\",\n \" def __init__(self,MAX_BOX):\\n\",\n \" super(Sampling,self).__init__()\\n\",\n \" self.max_box = MAX_BOX\\n\",\n \" \\n\",\n \" def forward(self,lamda):\\n\",\n \" \\n\",\n \" lamda = lamda.view(-1)\\n\",\n \" mask = T.zeros(lamda.shape[0] , self.max_box)\\n\",\n \" lamda = T.t(T.t(mask) + lamda)\\n\",\n \" mask = mask + T.arange(0,self.max_box,1)\\n\",\n \" e_lamda = T.exp(lamda)\\n\",\n \" lamda_x = lamda ** mask \\n\",\n \" fact = T.exp(T.lgamma(T.arange(0 , self.max_box)+1))\\n\",\n \" \\n\",\n \" # P = ((lambda ^ x)*e^(lamda)) / x! \\n\",\n \" probab = (lamda_x*e_lamda)/fact\\n\",\n \" sample = T.argmax(probab,dim=1)\\n\",\n \"\\n\",\n \" return sample \"\n ],\n \"execution_count\": null,\n \"outputs\": []\n },\n {\n \"cell_type\": \"code\",\n \"metadata\": {\n \"id\": \"2xi2UNH3Tl27\",\n \"execution\": {\n \"iopub.status.busy\": \"2021-07-22T12:08:54.457132Z\",\n \"iopub.execute_input\": \"2021-07-22T12:08:54.457527Z\",\n \"iopub.status.idle\": \"2021-07-22T12:08:54.473628Z\",\n \"shell.execute_reply.started\": \"2021-07-22T12:08:54.457492Z\",\n \"shell.execute_reply\": \"2021-07-22T12:08:54.472727Z\"\n },\n \"trusted\": true\n },\n \"source\": [\n \"class CountVAE(T.nn.Module):\\n\",\n \" \\n\",\n \" def __init__(self,n_class,max_box=9):\\n\",\n \" super(CountVAE,self).__init__()\\n\",\n \" \\n\",\n \" \\n\",\n \" self.encoder = Encoder()\\n\",\n \" self.prior = Prior()\\n\",\n \" self.decoder = Decoder(1)\\n\",\n \" self.embeder = Embeder(n_class)\\n\",\n \" self.loss = ELBOLoss() \\n\",\n \" self.rep = Reparamatrize()\\n\",\n \" self.n_class = n_class\\n\",\n \" self.pois = Sampling(max_box)\\n\",\n \" \\n\",\n \" def forward(self, inputs, isTrain = False):\\n\",\n \" \\n\",\n \" '''\\n\",\n \" isTrain(boolean) default False : defines whether data is to be treated as training data or testing\\n\",\n \" \\n\",\n \" if isTrain = True :\\n\",\n \" input must be a tuple with first value corresponding to label set and second corresponding to ground Truth\\n\",\n \" counts\\n\",\n \" else :\\n\",\n \" input must have label set\\n\",\n \" \\n\",\n \" '''\\n\",\n \" if isTrain==True:\\n\",\n \" \\n\",\n \" label_set , groundtruth_counts = inputs\\n\",\n \" Loss = 0\\n\",\n \" LL = 0\\n\",\n \" KL = 0\\n\",\n \" previous_counts = T.zeros_like(label_set)\\n\",\n \" \\n\",\n \" for i in range(self.n_class):\\n\",\n \" \\n\",\n \" current_label = T.zeros_like(previous_counts)\\n\",\n \" x_ = label_set[...,i]\\n\",\n \" current_label[...,i]= x_\\n\",\n \" z_ = groundtruth_counts[...,i].view(-1,1)\\n\",\n \" \\n\",\n \" # Generate Conditional Embedding\\n\",\n \" embedding = self.embeder([label_set, current_label, previous_counts])\\n\",\n \" \\n\",\n \" # Encoding To latet space\\n\",\n \" mu1, logvar1 = self.encoder([z_,embedding])\\n\",\n \" mu2, logvar2 = self.prior(embedding)\\n\",\n \" \\n\",\n \" # Reparamatrized Latent variable\\n\",\n \" z = self.rep([mu1,logvar1])\\n\",\n \"\\n\",\n \" # Decode from Latent space\\n\",\n \" decoded = self.decoder([embedding,z])\\n\",\n \" Closs, L_, kl_ = self.loss([mu1, logvar1, mu2, logvar2, decoded , z_])\\n\",\n \" \\n\",\n \" # Update Losses\\n\",\n \" Loss = Loss + Closs\\n\",\n \" LL = LL + L_\\n\",\n \" KL = KL + kl_\\n\",\n \" \\n\",\n \" decoded = T.exp(decoded)\\n\",\n \" \\n\",\n \" # Poisson Distributions with rate of Deoded\\n\",\n \" # q = self.pois(decoded)\\n\",\n \" q = Poisson(decoded).sample()\\n\",\n \" \\n\",\n \" # update Preivious Counts\\n\",\n \" previous_counts = previous_counts + current_label*(q.view(-1,1) + x_.view(-1,1))\\n\",\n \" \\n\",\n \" return Loss/self.n_class, KL/self.n_class, LL/self.n_class\\n\",\n \" \\n\",\n \" else:\\n\",\n \" \\n\",\n \" label_set = inputs\\n\",\n \" previous_counts = T.zeros_like(label_set)\\n\",\n \" \\n\",\n \" for i in range(self.n_class):\\n\",\n \"\\n\",\n \" current_label = T.zeros_like(previous_counts)\\n\",\n \" x_ = label_set[...,i]\\n\",\n \" current_label[...,i]= x_\\n\",\n \" \\n\",\n \" \\n\",\n \" # Generate Conditional Embedding\\n\",\n \" embedding = self.embeder([label_set, current_label, previous_counts])\\n\",\n \" \\n\",\n \" # Encoding To latet space\\n\",\n \" mu,logvar = self.prior(embedding)\\n\",\n \" \\n\",\n \" # Reparamatrized Latent variable\\n\",\n \" z = self.rep([mu,logvar])\\n\",\n \" \\n\",\n \" # Decode from Latent space\\n\",\n \" decoded = self.decoder([embedding,z])\\n\",\n \" decoded = T.exp(decoded)\\n\",\n \"\\n\",\n \" # Poisson Distributions with rate of Deoded\\n\",\n \" # q = self.pois(decoded)\\n\",\n \" q = Poisson(decoded).sample()\\n\",\n \" \\n\",\n \" # update Preivious Counts\\n\",\n \" previous_counts = previous_counts + current_label*(q.view(-1,1) + x_.view(-1,1))\\n\",\n \" \\n\",\n \" return previous_counts\"\n ],\n \"execution_count\": null,\n \"outputs\": []\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {\n \"id\": \"W4LsYqsnngjh\"\n },\n \"source\": [\n \"\\n\",\n \"# BboxVAE Model Architecture\\n\",\n \"\\n\",\n \"### Classes\\n\",\n \"1. Condtional Embedder\\n\",\n \"2. Encoder\\n\",\n \"3. Prior\\n\",\n \"4. Decoder\\n\",\n \"\\n\",\n \"### Loss\\n\",\n \"1. ELBO LOSS\\n\",\n \"\\n\",\n \"### Reparamatrize\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"metadata\": {\n \"id\": \"9-3b9UJxngji\",\n \"execution\": {\n \"iopub.status.busy\": \"2021-07-22T09:47:17.692836Z\",\n \"iopub.execute_input\": \"2021-07-22T09:47:17.693151Z\",\n \"iopub.status.idle\": \"2021-07-22T09:47:17.701175Z\",\n \"shell.execute_reply.started\": \"2021-07-22T09:47:17.693120Z\",\n \"shell.execute_reply\": \"2021-07-22T09:47:17.699927Z\"\n },\n \"trusted\": true\n },\n \"source\": [\n \"class EmbedBbox(T.nn.Module):\\n\",\n \" \\n\",\n \" def __init__(self,n_class):\\n\",\n \" super(EmbedBbox,self).__init__()\\n\",\n \" \\n\",\n \" self.fcb1 = fcblock(n_class)\\n\",\n \" self.fcb2 = fcblock(n_class)\\n\",\n \" self.seq1 = Sequential(\\n\",\n \" Linear(128,128),\\n\",\n \" ReLU()\\n\",\n \" )\\n\",\n \" \\n\",\n \" self.n_class = n_class\\n\",\n \" self.fc = Linear(128*3,128)\\n\",\n \" self.lstm = LSTM(n_class+4, hidden_size=128)\\n\",\n \"\\n\",\n \" def forward(self,inputs):\\n\",\n \" \\n\",\n \" in1,in2,in3 = inputs\\n\",\n \"\\n\",\n \" _ , (h_0 , c_0 ) = self.lstm(in3)\\n\",\n \" hn = h_0.view(-1, 128)\\n\",\n \" \\n\",\n \" in1 = self.fcb1(in1)\\n\",\n \" in2 = self.fcb2(in2)\\n\",\n \" in3 = self.seq1(hn)\\n\",\n \" \\n\",\n \" out = T.cat((in1,in2,in3),1)\\n\",\n \" out = self.fc(out)\\n\",\n \" \\n\",\n \" return out\"\n ],\n \"execution_count\": null,\n \"outputs\": []\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {\n \"id\": \"FfctLBtgoCLh\"\n },\n \"source\": [\n \"# Loss Function and Reparamatrization for BboxVAE\\n\",\n \"\\n\",\n \"## KL Divergence\\n\",\n \" Same as CountVAE\\n\",\n \"\\n\",\n \"## MSE\\n\",\n \" as reconstruction loss\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"metadata\": {\n \"id\": \"g9naXD3ln_1V\",\n \"execution\": {\n \"iopub.status.busy\": \"2021-07-22T09:47:18.720808Z\",\n \"iopub.execute_input\": \"2021-07-22T09:47:18.721123Z\",\n \"iopub.status.idle\": \"2021-07-22T09:47:18.730454Z\",\n \"shell.execute_reply.started\": \"2021-07-22T09:47:18.721091Z\",\n \"shell.execute_reply\": \"2021-07-22T09:47:18.729627Z\"\n },\n \"trusted\": true\n },\n \"source\": [\n \"class ELBOLoss_Bbox(T.nn.Module):\\n\",\n \" \\n\",\n \" def __init__(self):\\n\",\n \" super(ELBOLoss_Bbox,self).__init__()\\n\",\n \" \\n\",\n \" def forward(self,inputs):\\n\",\n \" mu1,logvar1,mu2,logvar2, xp , yp = inputs\\n\",\n \" \\n\",\n \" ''' KL Divergence '''\\n\",\n \" kl = 0.5 * T.sum((logvar2 - logvar1) - 1 + (logvar1.exp() + (mu2 - mu1).pow(2) )/logvar2.exp() , dim = -1 ).mean()\\n\",\n \" \\n\",\n \" ''' Multivariate Guassian Likelihood '''\\n\",\n \" mse = T.nn.MSELoss()(xp,yp)\\n\",\n \" loss = mse + kl\\n\",\n \" \\n\",\n \" return loss, kl,mse\\n\",\n \"\\n\",\n \"\\n\",\n \"class Reparamatrize(T.nn.Module):\\n\",\n \" \\n\",\n \" def __init__(self):\\n\",\n \" super(Reparamatrize,self).__init__()\\n\",\n \" \\n\",\n \" def forward(self,inputs):\\n\",\n \" \\n\",\n \" mu , logvar = inputs\\n\",\n \" std = T.exp(logvar/2)\\n\",\n \" eps = T.rand_like(std)\\n\",\n \"\\n\",\n \" return eps*std + mu\\n\",\n \" \\n\",\n \"\\n\",\n \"class ReparamatrizeMulti(T.nn.Module):\\n\",\n \" \\n\",\n \" def __init__(self):\\n\",\n \" super(ReparamatrizeMulti,self).__init__()\\n\",\n \" \\n\",\n \" def forward(self,inputs):\\n\",\n \" \\n\",\n \" mu = inputs\\n\",\n \" std = (T.ones_like(mu)*0.02)\\n\",\n \" eps = T.rand_like(std)\\n\",\n \" \\n\",\n \" return eps*std + mu\"\n ],\n \"execution_count\": null,\n \"outputs\": []\n },\n {\n \"cell_type\": \"code\",\n \"metadata\": {\n \"id\": \"2F1cTxFingji\",\n \"execution\": {\n \"iopub.status.busy\": \"2021-07-22T09:47:19.239020Z\",\n \"iopub.execute_input\": \"2021-07-22T09:47:19.239349Z\",\n \"iopub.status.idle\": \"2021-07-22T09:47:19.257000Z\",\n \"shell.execute_reply.started\": \"2021-07-22T09:47:19.239314Z\",\n \"shell.execute_reply\": \"2021-07-22T09:47:19.255889Z\"\n },\n \"trusted\": true\n },\n \"source\": [\n \"class BboxVAE(T.nn.Module):\\n\",\n \" def __init__(self,n_class,n_dim,max_box,latent_dim=32):\\n\",\n \"\\n\",\n \" super(BboxVAE,self).__init__()\\n\",\n \" \\n\",\n \" self.embeder = EmbedBbox(n_class)\\n\",\n \" self.encoder = Encoder(n_dim,latent_dim=latent_dim)\\n\",\n \" self.decoder = Decoder(n_dim,latent_dim=latent_dim)\\n\",\n \" self.prior = Prior(latent_dim=latent_dim)\\n\",\n \" self.loss = ELBOLoss_Bbox()\\n\",\n \" self.rep = Reparamatrize()\\n\",\n \" self.n_dim = n_dim\\n\",\n \" self.n_class = n_class\\n\",\n \" self.rep_mul = ReparamatrizeMulti()\\n\",\n \" self.max_box = max_box\\n\",\n \"\\n\",\n \"\\n\",\n \" def forward(self,inputs,isTrain=True):\\n\",\n \" if isTrain==True :\\n\",\n \" BoxCounts, GTBBox , BoxLabel= inputs\\n\",\n \" los = 0\\n\",\n \" kl1 = 0\\n\",\n \" ll1 = 0\\n\",\n \" for i in range(self.max_box):\\n\",\n \" if i==0:\\n\",\n \" PrevLabel = T.zeros((1 , *BoxLabel[... ,i,:].shape)) \\n\",\n \" PrevBox = T.zeros((1 , *GTBBox[...,i,:].shape))\\n\",\n \" \\n\",\n \"\\n\",\n \" GroundTruth = GTBBox[... , i ,:].view(-1,self.n_dim)\\n\",\n \" \\n\",\n \" CurrentLabel = BoxLabel[... , i ,:].view(-1,self.n_class)\\n\",\n \" \\n\",\n \" Embedding = self.embeder([BoxCounts,CurrentLabel,T.cat([PrevLabel,PrevBox] , dim = 2)])\\n\",\n \"\\n\",\n \" mu1 , logvar1 = self.encoder([GroundTruth,Embedding])\\n\",\n \" mu2 , logvar2 = self.prior(Embedding)\\n\",\n \" z1 = self.rep([mu1,logvar1])\\n\",\n \" z2 = self.rep([mu2,logvar2])\\n\",\n \" \\n\",\n \" Mu = self.decoder([Embedding,z1])\\n\",\n \" BBox = self.rep_mul(Mu)\\n\",\n \" CLoss, kl_tot , ll_tot = self.loss([mu1,logvar1,mu2,logvar2, BBox , GroundTruth])\\n\",\n \"\\n\",\n \" los = los + CLoss/self.max_box\\n\",\n \" kl1 = kl1 + kl_tot/self.max_box\\n\",\n \" ll1 = ll1 + ll_tot/self.max_box\\n\",\n \" \\n\",\n \" PrevBox = T.cat([PrevBox ,T.unsqueeze(GroundTruth,0)])\\n\",\n \" PrevLabel = T.cat([PrevLabel , T.unsqueeze(CurrentLabel,0)])\\n\",\n \"\\n\",\n \"\\n\",\n \" return los , kl1 , ll1\\n\",\n \" else:\\n\",\n \" BoxCounts, BoxLabel= inputs\\n\",\n \" BBoxes = []\\n\",\n \" for i in range(self.max_box):\\n\",\n \" if i==0:\\n\",\n \" PrevLabel = T.zeros((1 , *BoxLabel[... ,i,:].shape)) \\n\",\n \" PrevBox = T.zeros((1 , BoxLabel.shape[0] , 4))\\n\",\n \"\\n\",\n \" CurrentLabel = BoxLabel[... , i ,:].view(-1,self.n_class)\\n\",\n \" Embedding = self.embeder([BoxCounts,CurrentLabel,T.cat([PrevLabel,PrevBox] , dim = 2)])\\n\",\n \" \\n\",\n \" mu , logvar = self.prior(Embedding)\\n\",\n \" \\n\",\n \" z = self.rep([mu,logvar])\\n\",\n \" \\n\",\n \" Mu = self.decoder([Embedding,z])\\n\",\n \" \\n\",\n \" BBox = self.rep_mul(Mu)\\n\",\n \" \\n\",\n \" PrevBox = T.cat([PrevBox ,T.unsqueeze(BBox,0)])\\n\",\n \" PrevLabel = T.cat([PrevLabel , T.unsqueeze(CurrentLabel,0)])\\n\",\n \" BBoxes.append(BBox.t())\\n\",\n \" BBoxes =T.stack(BBoxes)\\n\",\n \" return BBoxes\"\n ],\n \"execution_count\": null,\n \"outputs\": []\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {\n \"id\": \"ZVDs9sdm9E06\"\n },\n \"source\": [\n \"# Layout VAE\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"metadata\": {\n \"id\": \"rl17zTEq9D0S\"\n },\n \"source\": [\n \"class LayoutVAE(T.nn.Module):\\n\",\n \"\\n\",\n \" def __init__(self, n_class = 6, max_box = 9,bboxvae_latent_dim = 32,bboxvae_lr=1e-4,countvae_lr=1e-6):\\n\",\n \" '''\\n\",\n \" ** Layout VAE **\\n\",\n \" * https://arxiv.org/abs/1907.10719\\n\",\n \" '''\\n\",\n \" super(LayoutVAE,self).__init__()\\n\",\n \"\\n\",\n \" self.max_box = max_box\\n\",\n \" self.n_class = n_class\\n\",\n \" self.lr_bvae = bboxvae_lr\\n\",\n \" self.lr_cvae = countvae_lr\\n\",\n \" self.countvae = CountVAE(n_class)\\n\",\n \" self.bboxvae = BboxVAE(n_class,4,max_box,bboxvae_latent_dim)\\n\",\n \" self.is_cvae_trained = 0\\n\",\n \" self.is_bvae_trained = 0\\n\",\n \"\\n\",\n \" def forward(self,input):\\n\",\n \" '''\\n\",\n \" Takes only Labels Set as input\\n\",\n \" Label Set : it is a vector of size n_class and contains 1 if correspinding class is present\\n\",\n \" '''\\n\",\n \" if self.is_cvae_trained == 0:\\n\",\n \" print(\\\"[Warning] Count VAE is Not Trained !!\\\")\\n\",\n \"\\n\",\n \" if self.is_bvae_trained == 0:\\n\",\n \" print(\\\"[Warning] Bbox VAE is Not Trained !!\\\")\\n\",\n \"\\n\",\n \" label_set = input\\n\",\n \" pred_class_counts = self.countvae(label_set , isTrain=False)\\n\",\n \" \\n\",\n \" # Normalize classiction between [0 , max_box]\\n\",\n \" pred_class_counts = T.floor ( self.max_box*(pred_class_counts / T.sum(pred_class_counts , dim = 1 ).view(-1,1)) )\\n\",\n \" \\n\",\n \" # Extra boxes which are not be predicted\\n\",\n \" # Their counts are set in first class\\n\",\n \" for class_count in pred_class_counts:\\n\",\n \" if(T.sum(class_count) < self.max_box):\\n\",\n \" class_count[0] = self.max_box - T.sum(class_count)\\n\",\n \"\\n\",\n \" class_labels = T.zeros(len(label_set) , self.max_box, self.n_class)\\n\",\n \"\\n\",\n \" for i in range(len(pred_class_counts)):\\n\",\n \" l = 0\\n\",\n \" for j in range(self.n_class):\\n\",\n \" for k in range(int(pred_class_counts[i][self.n_class-j-1])):\\n\",\n \" class_labels[i][l][self.n_class-j-1] = 1;\\n\",\n \" l+=1\\n\",\n \"\\n\",\n \" pred_box = self.bboxvae([ pred_class_counts, class_labels], isTrain=False)\\n\",\n \" pred_box = pred_box.permute(2,0,1)\\n\",\n \" class_info = T.unsqueeze(T.argmax(class_labels ,dim=2),dim=2)\\n\",\n \" predictions = T.cat([class_info,pred_box],dim = 2)\\n\",\n \"\\n\",\n \" for i in range(len(predictions)):\\n\",\n \" for j in range(len(predictions[i])):\\n\",\n \" if predictions[i][j][0]==0:\\n\",\n \" predictions[i][j]*=0\\n\",\n \" \\n\",\n \" self.predictions = predictions\\n\",\n \" self.pred_class_counts = pred_class_counts\\n\",\n \"\\n\",\n \" return predictions\\n\",\n \"\\n\",\n \" def load_data(self, path, frac = 0.5, train_test_split = 0.1):\\n\",\n \" '''\\n\",\n \" Loads data from npy file\\n\",\n \" path string containig path to data\\n\",\n \" frac defines the fraction of data to load\\n\",\n \"\\n\",\n \" '''\\n\",\n \" try : \\n\",\n \" Data = np.load(DATA_PATH)\\n\",\n \" # Sortind Data in proper order\\n\",\n \" order = np.argsort(Data[:,:,0])\\n\",\n \" for i in range(len(Data)):\\n\",\n \" Data[i] = Data[i][order[i][::-1]]\\n\",\n \" np.random.shuffle(Data)\\n\",\n \"\\n\",\n \" data_size = int(frac*len(Data))\\n\",\n \" test_size = int(train_test_split*data_size)\\n\",\n \" Data = T.tensor(Data[0:data_size]).float()\\n\",\n \" test_data = Data[0:test_size]\\n\",\n \" Data = Data[test_size:]\\n\",\n \"\\n\",\n \" # Prepare Data\\n\",\n \" self.class_labels = Data[...,4:]\\n\",\n \" self.class_counts = T.sum(Data[...,4:], dim = 1)\\n\",\n \" self.b_boxes = Data[...,0:4]\\n\",\n \" self.label_set = (self.class_counts !=0) + 0.0\\n\",\n \"\\n\",\n \" # Test Data\\n\",\n \" self.test_class_labels = test_data[...,4:]\\n\",\n \" self.test_class_counts = T.sum(test_data[...,4:], dim = 1)\\n\",\n \" self.test_b_boxes = test_data[...,0:4]\\n\",\n \" self.test_label_set = (self.test_class_counts !=0) + 0.0\\n\",\n \"\\n\",\n \" print(\\\"[Success] Data Loaded Succesfully\\\")\\n\",\n \"\\n\",\n \" except: \\n\",\n \" print(\\\"[Failed] Data Loading Failed\\\\n please check path\\\")\\n\",\n \" \\n\",\n \" def train(self, optim, train_mode = 'bboxvae', epochs = 100, bsize = 256 , validation_split = 0.1):\\n\",\n \" '''\\n\",\n \" * train_mode (str , default bboxvae) : Two optons\\n\",\n \" 1. if train_mode is bboxvae, BBoxVAE model will be trained and data \\n\",\n \" will be loaded accordingly\\n\",\n \" 2. if train_mode is countvae, CountVAE model will be trained and data \\n\",\n \" will be loaded accordingly\\n\",\n \" * epochs (int , default 100 ) : number of epochs training should run\\n\",\n \" * bsize(int default 256) : Batch Size\\n\",\n \" * validation_split(float default 0.1) : should be between between 0 and 1\\n\",\n \" 1 . it defines the size of validation data \\n\",\n \"\\n\",\n \" '''\\n\",\n \" # Create validation Split\\n\",\n \" total_examples = len(self.class_counts)\\n\",\n \" val_size = int(total_examples*validation_split)\\n\",\n \" \\n\",\n \" losses = dict()\\n\",\n \" train_data = []\\n\",\n \" if train_mode == 'countvae':\\n\",\n \" model = self.countvae\\n\",\n \" train_data = [self.label_set, self.class_counts]\\n\",\n \" else :\\n\",\n \" model = self.bboxvae\\n\",\n \" train_data = [self.class_counts, self.b_boxes, self.class_labels]\\n\",\n \"\\n\",\n \" # Validation Data \\n\",\n \" val_data = []\\n\",\n \" for x in train_data:\\n\",\n \" val_data.append(x[:val_size])\\n\",\n \"\\n\",\n \" # Train data\\n\",\n \" for i in range(len(train_data)):\\n\",\n \" train_data[i] = train_data[i][val_size:]\\n\",\n \"\\n\",\n \"\\n\",\n \" # find the number of batches\\n\",\n \" batches = len(train_data[0])//bsize\\n\",\n \" second_loss = 'mse'\\n\",\n \" if train_mode == 'countvae':\\n\",\n \" second_loss = 'poisson_nll'\\n\",\n \" \\n\",\n \" # Dictionary to keep track of model statistics\\n\",\n \" losses = {'epoch':-1, \\n\",\n \" 'batch':0,\\n\",\n \" 'lr' : 0,\\n\",\n \" 'loss':0,\\n\",\n \" 'kl_div_loss':0,\\n\",\n \" second_loss+'_loss':0,\\n\",\n \" 'val_loss':0,\\n\",\n \" 'val_kl_div_loss':0,\\n\",\n \" 'val_'+second_loss+'_loss':0\\n\",\n \" }\\n\",\n \"\\n\",\n \" history = pd.DataFrame(losses ,index = [0])\\n\",\n \" index = 1\\n\",\n \"\\n\",\n \" for ep in range(epochs):\\n\",\n \"\\n\",\n \" # if train_mode=='countvae':\\n\",\n \" # self.countvae_pred_grpah(epoch = ep,path = CVAE_PATH)\\n\",\n \"\\n\",\n \" print(f'Epoch[{ep+1}/{epochs}]')\\n\",\n \" for batch in range(batches):\\n\",\n \"\\n\",\n \" # Get Current batch\\n\",\n \" b = []\\n\",\n \" for x in train_data:\\n\",\n \" b.append(x[batch*bsize : (batch+1)*bsize])\\n\",\n \"\\n\",\n \" optim.zero_grad()\\n\",\n \"\\n\",\n \" # Train Step\\n\",\n \" loss, kl_, l_ = model(b,isTrain = True)\\n\",\n \" \\n\",\n \" # Validation Step\\n\",\n \" val_loss, val_kl_, val_l_ = model(val_data, isTrain = True)\\n\",\n \"\\n\",\n \"\\n\",\n \" # Save Statistics\\n\",\n \" losses['epoch'] = ep\\n\",\n \" losses['batch'] = batch\\n\",\n \" losses['lr'] = optim.param_groups[0]['lr']\\n\",\n \"\\n\",\n \" loss_list = [loss, kl_, l_ , val_loss , val_kl_ , val_l_]\\n\",\n \"\\n\",\n \" for i in range(6):\\n\",\n \" losses[list(losses.keys())[3+i]] = loss_list[i].cpu().clone().detach().numpy()\\n\",\n \" pass\\n\",\n \" \\n\",\n \" losses_df = pd.DataFrame(losses , index=[index])\\n\",\n \" history = pd.concat([history,losses_df])\\n\",\n \" index+=1\\n\",\n \"\\n\",\n \" # Backpropogation step and updating weights\\n\",\n \" loss.backward()\\n\",\n \" optim.step()\\n\",\n \" print('\\\\r Batch: {}/{} - loss : {} - val_loss : {} - val_{} : {}'.format(batch+1,batches,\\n\",\n \" losses_df['loss'][index-1],\\n\",\n \" losses_df['val_loss'][index-1],\\n\",\n \" second_loss,\\n\",\n \" losses_df['val_'+second_loss+'_loss'][index-1]),\\n\",\n \" end=\\\"\\\")\\n\",\n \" print(\\\"\\\\n\\\")\\n\",\n \" print('[Success] Finished Training')\\n\",\n \" return history\\n\",\n \"\\n\",\n \" def load_countvae_weights(self,path):\\n\",\n \" try :\\n\",\n \" self.countvae = T.load(path)\\n\",\n \" self.is_cvae_trained=1\\n\",\n \" print('[Success] Loaded Successfully')\\n\",\n \" except:\\n\",\n \" print('[Failed] Load Failed')\\n\",\n \"\\n\",\n \" def load_bboxvae_weights(self,path):\\n\",\n \" try :\\n\",\n \" self.bboxvae = T.load(path)\\n\",\n \" self.is_bvae_trained=1\\n\",\n \" print('[Success] Loaded Successfully')\\n\",\n \" except:\\n\",\n \" print('[Failed] Load Failed')\\n\",\n \"\\n\",\n \" def train_bboxvae(self,epochs=30, bsize=256, validation_split=0.1, optim=None): \\n\",\n \" if optim == None:\\n\",\n \" optim = T.optim.Adam(self.bboxvae.parameters(),lr=self.lr_bvae)\\n\",\n \"\\n\",\n \" # Start Training\\n\",\n \" history = self.train(optim = optim,\\n\",\n \" train_mode = 'bboxvae',\\n\",\n \" epochs = epochs,\\n\",\n \" bsize = bsize,\\n\",\n \" validation_split = validation_split\\n\",\n \" )\\n\",\n \" self.is_bvae_trained = 1\\n\",\n \" self.bvae_history = history[history.columns][1:]\\n\",\n \" return self.bvae_history\\n\",\n \"\\n\",\n \" def train_countvae(self,epochs=30, bsize=256, validation_split=0.1, optim=None):\\n\",\n \" \\n\",\n \" if optim == None:\\n\",\n \" optim = T.optim.Adam(self.countvae.parameters(),lr=self.lr_cvae)\\n\",\n \"\\n\",\n \" # Start Training\\n\",\n \" history = self.train(optim = optim,\\n\",\n \" train_mode = 'countvae',\\n\",\n \" epochs = epochs,\\n\",\n \" bsize = bsize,\\n\",\n \" validation_split = validation_split\\n\",\n \" )\\n\",\n \" self.is_cvae_trained = 1\\n\",\n \" self.cvae_history = history[history.columns][1:]\\n\",\n \" return self.cvae_history\\n\",\n \"\\n\",\n \" def pred_countvae(self,data=None):\\n\",\n \" '''\\n\",\n \" * Functions is used for for predcting from CountVAE\\n\",\n \" given label_set\\n\",\n \" * if data is None than label set from loaded data \\n\",\n \" are used for predictions.\\n\",\n \" '''\\n\",\n \"\\n\",\n \" if self.is_cvae_trained == 0:\\n\",\n \" print(\\\"[Warning] Count VAE is Not Trained !!\\\")\\n\",\n \" if data == None :\\n\",\n \" data = self.test_label_set\\n\",\n \" return self.countvae(data , isTrain=False)\\n\",\n \"\\n\",\n \" def pred_bboxvae(self,Data=None):\\n\",\n \"\\n\",\n \" '''\\n\",\n \" * Functions is used for for predcting from BboxVAE\\n\",\n \" given class_counts and class labels\\n\",\n \" * if data is None than class counts and class labels from loaded data \\n\",\n \" are used for predictions.\\n\",\n \" '''\\n\",\n \"\\n\",\n \" if self.is_bvae_trained == 0:\\n\",\n \" print(\\\"[Warning] Bbox VAE is Not Trained !!\\\")\\n\",\n \"\\n\",\n \" if Data == None :\\n\",\n \" Data = [self.test_class_counts,self.test_class_labels]\\n\",\n \"\\n\",\n \" batches = len(Data[0])//64\\n\",\n \"\\n\",\n \" for b in range(batches):\\n\",\n \"\\n\",\n \" # Get data in batch\\n\",\n \" data = [self.test_class_counts[b*64 : (b+1)*64],\\n\",\n \" self.test_class_labels[b*64 : (b+1)*64]] \\n\",\n \"\\n\",\n \" # Predict\\n\",\n \" pred = self.bboxvae(data, isTrain=False)\\n\",\n \" pred = pred.permute(2,0,1)\\n\",\n \"\\n\",\n \" # cxywh format\\n\",\n \" class_info = T.unsqueeze(T.argmax(data[1] ,dim=2),dim=2)\\n\",\n \" pred = T.cat([class_info,pred],dim = 2)\\n\",\n \"\\n\",\n \"\\n\",\n \" for i in range(len(pred)):\\n\",\n \" for j in range(len(pred[i])):\\n\",\n \" if pred[i][j][0]==0:\\n\",\n \" pred[i][j] *= 0\\n\",\n \"\\n\",\n \" if b > 0:\\n\",\n \" predictions = T.cat([predictions,pred],dim=0)\\n\",\n \" else:\\n\",\n \" predictions = pred\\n\",\n \" class_info =T.argmax(self.test_class_labels[0:64*batches] ,dim=2)\\n\",\n \" class_info = T.unsqueeze(class_info,dim=2)\\n\",\n \" gt = T.cat([class_info,self.test_b_boxes[0:64*batches]],dim = 2)\\n\",\n \" return predictions, gt\\n\",\n \" \\n\",\n \" def countvae_pred_grpah(self,path,epoch = 0):\\n\",\n \" pred_cvae = self.pred_countvae()\\n\",\n \" pred_cvae = T.sum(pred_cvae,dim=0)\\n\",\n \" pred_cvae = pred_cvae/T.sum(pred_cvae)\\n\",\n \" pred_cvae = pred_cvae.to('cpu').clone().detach().numpy()\\n\",\n \"\\n\",\n \" gt_cvae = T.sum(self.class_counts,dim=0)\\n\",\n \" gt_cvae = gt_cvae/T.sum(gt_cvae)\\n\",\n \" gt_cvae = gt_cvae.to('cpu').clone().detach().numpy()\\n\",\n \"\\n\",\n \" fig = plt.figure(figsize=(5 ,4), dpi=100 ,facecolor=(0,0,0))\\n\",\n \" ax = fig.add_subplot()\\n\",\n \" ax.plot(gt_cvae , 'red',marker = 'o', label = 'Ground Truth',linewidth=4)\\n\",\n \" ax.plot(pred_cvae,'blue',marker ='o',label = \\\"Predicted\\\" ,linewidth=4)\\n\",\n \" ax.legend()\\n\",\n \" ax.set_title('Ground Truth vs Predicted Distribution\\\\n Epoch = '+str(epoch))\\n\",\n \" ax.set_xlabel('Classes')\\n\",\n \" ax.set_xticks([0,1,2,3,4,5])\\n\",\n \" ax.set_xticklabels(class_names)\\n\",\n \"\\n\",\n \" plt.savefig(path+\\\"cvae-train-ep-\\\"+str(epoch)+\\\".png\\\",facecolor=(0,0,0))\\n\",\n \" plt.close()\\n\",\n \"\\n\",\n \"\\n\",\n \" def convert_to_cxywh(self,data):\\n\",\n \" \\n\",\n \" bboxes = data[...,0:4]\\n\",\n \" labels = data[...,4: ]\\n\",\n \" class_info = T.unsqueeze(T.argmax(labels ,dim=2),dim=2)\\n\",\n \" cxywh = T.cat([class_info,bboxes],dim = 2)\\n\",\n \" return cxywh\\n\",\n \"\\n\",\n \" def save_model(self,path):\\n\",\n \" \\n\",\n \" T.save(self.countvae,path+'countvae.h5')\\n\",\n \" T.save(self.bboxvae,path+'bboxvae.h5')\\n\",\n \" T.save(self,path+'selef.h5')\\n\",\n \" print('[Success] Saved Successfully')\\n\",\n \"\\n\",\n \" def save_history(self,path):\\n\",\n \"\\n\",\n \" self.cvae_history.to_csv(path+'cvae-history.csv',index=False)\\n\",\n \" self.bvae_history.to_csv(path+'bvae-history.csv',index=False)\\n\",\n \" print('[Success] Saved Successfully')\\n\",\n \"\\n\",\n \" \\n\"\n ],\n \"execution_count\": null,\n \"outputs\": []\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {\n \"id\": \"DpwhFKH_aFaF\"\n },\n \"source\": [\n \"# Plotting Functions\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"metadata\": {\n \"id\": \"BOjuR2gnSDvD\"\n },\n \"source\": [\n \"def plot_history(history,title = 'Training Statistics', path =\\\"\\\"):\\n\",\n \" height = 12\\n\",\n \" width = 9\\n\",\n \" fig = plt.figure(figsize=(width,height), dpi=100 ,facecolor=(0,0,0))\\n\",\n \" height_ratio = [0.25,1,1,1]\\n\",\n \" grid = plt.GridSpec(4,2,\\n\",\n \" hspace=0.3,wspace=0.2,\\n\",\n \" height_ratios =height_ratio,\\n\",\n \" left=0.02,right=0.98,top=0.98,bottom=0.02\\n\",\n \" )\\n\",\n \" index = 0\\n\",\n \" ax = fig.add_subplot(grid[index : index+2])\\n\",\n \" index+=2\\n\",\n \" ax.text(x = 0.3 ,y = 0.5 ,s = title,fontsize=30)\\n\",\n \" ax.invert_yaxis()\\n\",\n \" ax.axis('off')\\n\",\n \" colors = ['red','blue','green']\\n\",\n \" for i in range(3):\\n\",\n \"\\n\",\n \" ax = fig.add_subplot(grid[index])\\n\",\n \" ax.plot(history[history.columns[i+3]],colors[i])\\n\",\n \" index+=1\\n\",\n \" ax.set_facecolor((0,0,0))\\n\",\n \" ax.set_title(history.columns[i+3])\\n\",\n \" ax = fig.add_subplot(grid[index])\\n\",\n \" ax.plot(history[history.columns[i+6]],colors[i])\\n\",\n \" ax.set_title(history.columns[i+6])\\n\",\n \" index+=1\\n\",\n \" ax.set_facecolor((0,0,0))\\n\",\n \" plt.savefig(path, facecolor=(0,0,0))\"\n ],\n \"execution_count\": null,\n \"outputs\": []\n },\n {\n \"cell_type\": \"code\",\n \"metadata\": {\n \"id\": \"mBxvgs0-SUhL\"\n },\n \"source\": [\n \"def generate_colors(class_names = None,n_class=6):\\n\",\n \" cmap = [\\\"\\\",\\\"#dc143c\\\",\\\"#ffff00\\\",\\\"#00ff00\\\",\\\"#ff00ff\\\",\\\"#1e90ff\\\",\\\"#fff5ee\\\",\\n\",\n \" \\\"#00ffff\\\",\\\"#8b008b\\\",\\\"#ff4500\\\",\\\"#8b4513\\\",\\\"#808000\\\",\\\"#483d8b\\\",\\n\",\n \" \\\"#008000\\\",\\\"#000080\\\",\\\"#9acd32\\\",\\\"#ffa500\\\",\\\"#ba55d3\\\",\\\"#00fa9a\\\",\\n\",\n \" \\\"#dc143c\\\",\\\"#0000ff\\\",\\\"#f08080\\\",\\\"#f0e68c\\\",\\\"#dda0dd\\\",\\\"#ff1493\\\"]\\n\",\n \" \\n\",\n \" colors = dict()\\n\",\n \"\\n\",\n \" if class_names == None:\\n\",\n \" class_names = []\\n\",\n \" for i in range(n_class):\\n\",\n \" class_names.append('class'+str(i+1))\\n\",\n \" \\n\",\n \" for i in range(n_class):\\n\",\n \" colors[class_names[i]] = cmap[i]\\n\",\n \"\\n\",\n \" return colors\\n\",\n \"\\n\",\n \"def plot_layouts(data,colors,class_names,title=\\\"Random Predictions\\\", path=\\\"\\\"):\\n\",\n \" '''\\n\",\n \" data in cxywh format\\n\",\n \" '''\\n\",\n \" height = 15\\n\",\n \" width = 9\\n\",\n \" fig = plt.figure(figsize=(width,height), dpi=100 ,facecolor=(0,0,0))\\n\",\n \" height_ratio = [0.5,0.25,1,1,1,1]\\n\",\n \" grid = plt.GridSpec(6,4,\\n\",\n \" hspace=0.05,wspace=0.05,\\n\",\n \" height_ratios =height_ratio,\\n\",\n \" left=0.02,right=0.98,top=0.98,bottom=0.02\\n\",\n \" )\\n\",\n \" index = 0\\n\",\n \"\\n\",\n \"\\n\",\n \" ax = fig.add_subplot(grid[index : index+4])\\n\",\n \" index+=4\\n\",\n \" ax.text(x = 0.2 ,y = 0.5 ,s = title,fontsize=30)\\n\",\n \" ax.axis('off')\\n\",\n \" legend = []\\n\",\n \" ax = fig.add_subplot(grid[index : index+4])\\n\",\n \" index += 4\\n\",\n \" \\n\",\n \" for i in range(1,6):\\n\",\n \" legend.append(Patch(facecolor=colors[class_names[i]]+\\\"40\\\",\\n\",\n \" edgecolor=colors[class_names[i]],\\n\",\n \" label= class_names[i]))\\n\",\n \" \\n\",\n \" ax.legend(handles=legend, ncol=3,loc=8, fontsize=25, facecolor=(0,0,0))\\n\",\n \" ax.axis('off')\\n\",\n \"\\n\",\n \" for i in range(16):\\n\",\n \" ax = fig.add_subplot(grid[index])\\n\",\n \" index += 1\\n\",\n \" \\n\",\n \" data = pred[i]\\n\",\n \" rect1 = patches.Rectangle((0,0),180,240)\\n\",\n \" rect1.set_color((0,0,0,1))\\n\",\n \" ax.add_patch(rect1)\\n\",\n \" for box in data:\\n\",\n \"\\n\",\n \" c,x,y,w,h = box\\n\",\n \" if c==0:\\n\",\n \" continue\\n\",\n \" x = x*180\\n\",\n \" y = y*240\\n\",\n \" w = w*180\\n\",\n \" h = h*240\\n\",\n \" rect = patches.Rectangle((x,y),w,h,linewidth=2)\\n\",\n \" rect.set_color(colors[class_names[int(c)]]+\\\"72\\\")\\n\",\n \" rect.set_linestyle('-')\\n\",\n \" rect.set_edgecolor(colors[class_names[int(c)]])\\n\",\n \" ax.add_patch(rect)\\n\",\n \" ax.plot()\\n\",\n \" ax.set_facecolor((0,0,0))\\n\",\n \" for spine in ax.spines.values():\\n\",\n \" spine.set_edgecolor('green')\\n\",\n \" spine.set_linewidth(2)\\n\",\n \" ax.invert_yaxis()\\n\",\n \" ax.set_xticks([])\\n\",\n \" ax.set_yticks([])\\n\",\n \" plt.savefig(path, facecolor=(0,0,0))\\n\",\n \" plt.show()\\n\",\n \" plt.close()\"\n ],\n \"execution_count\": null,\n \"outputs\": []\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {\n \"id\": \"BRUduBg28naW\"\n },\n \"source\": [\n \"# Training\\n\",\n \"\\n\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"metadata\": {\n \"id\": \"eiWxIjvZ80dH\"\n },\n \"source\": [\n \"layoutvae = LayoutVAE()\\n\",\n \"layoutvae.load_data(DATA_PATH, frac = 0.5)\"\n ],\n \"execution_count\": null,\n \"outputs\": []\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {\n \"id\": \"muczxUPh8vfp\"\n },\n \"source\": [\n \"## Countvae\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"metadata\": {\n \"id\": \"xc9TZ6DfTCrY\"\n },\n \"source\": [\n \"# layoutvae.load_countvae_weights(path = SAVE_MODEL_PATH + \\\"countvae.h5\\\")\\n\",\n \"layoutvae.train_countvae(bsize = 512, epochs=100,validation_split=0.1)\\n\",\n \"plot_history(layoutvae.cvae_history,\\n\",\n \" title=\\\"CountVAE Training\\\",\\n\",\n \" path = SAVE_LOG_PATH+\\\"Cvae-train.svg\\\"\\n\",\n \" )\\n\",\n \"\\n\",\n \"def countvae_pred_grpah(self,path):\\n\",\n \" pred_cvae = self.pred_countvae()\\n\",\n \" pred_cvae = T.sum(pred_cvae,dim=0)\\n\",\n \" pred_cvae = pred_cvae/T.sum(pred_cvae)\\n\",\n \" pred_cvae = pred_cvae.to('cpu').clone().detach().numpy()\\n\",\n \"\\n\",\n \" gt_cvae = T.sum(self.class_counts,dim=0)\\n\",\n \" gt_cvae = gt_cvae/T.sum(gt_cvae)\\n\",\n \" gt_cvae = gt_cvae.to('cpu').clone().detach().numpy()\\n\",\n \"\\n\",\n \" fig = plt.figure(figsize=(5 ,4), dpi=100 ,facecolor=(0,0,0))\\n\",\n \" ax = fig.add_subplot()\\n\",\n \" ax.plot(gt_cvae , 'red',marker = 'o', label = 'Ground Truth',linewidth=4)\\n\",\n \" ax.plot(pred_cvae,'blue',marker ='o',label = \\\"Predicted\\\" ,linewidth=4)\\n\",\n \" ax.legend()\\n\",\n \" ax.set_title('Ground Truth vs Predicted Distribution')\\n\",\n \" ax.set_xlabel('Classes')\\n\",\n \" ax.set_xticks([0,1,2,3,4,5])\\n\",\n \" ax.set_xticklabels(class_names)\\n\",\n \"\\n\",\n \" plt.savefig(path+\\\"cvae-train.png\\\",facecolor=(0,0,0))\\n\",\n \" plt.close()\"\n ],\n \"execution_count\": null,\n \"outputs\": []\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {\n \"id\": \"mYKbkK-GOn24\"\n },\n \"source\": [\n \"## Bbox VAE\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"metadata\": {\n \"id\": \"g-yqu2BBDb9A\"\n },\n \"source\": [\n \"# layoutvae.load_bboxvae_weights(path = SAVE_MODEL_PATH + \\\"bboxvae.h5\\\")\\n\",\n \"history_df = layoutvae.train_bboxvae(bsize = 256, epochs = 150, validation_split = 0.1)\\n\",\n \"preds,gt = layoutvae.pred_bboxvae()\\n\",\n \"plot_history(layoutvae.vae_history,\\n\",\n \" title=\\\"BBoxVAE Training\\\",\\n\",\n \" path = SAVE_LOG_PATH+\\\"Bvae-train.svg\\\"\\n\",\n \" )\"\n ],\n \"execution_count\": null,\n \"outputs\": []\n },\n {\n \"cell_type\": \"code\",\n \"metadata\": {\n \"id\": \"-9pIjabEZiTW\"\n },\n \"source\": [\n \"class_names = ['None' , 'Text' , 'Title' , 'List' , 'Table' ,'Figure']\\n\",\n \"colors = generate_colors(n_class=6 , class_names=class_names)\\n\",\n \"\\n\",\n \"preds = layoutvae.pred_bboxvae()\\n\",\n \"for i in range(2):\\n\",\n \" plot_layouts(data = predd[i*16:(i+1)*16],\\n\",\n \" colors=colors,\\n\",\n \" class_names=class_names,\\n\",\n \" path=SAVE_OUTPUT_PATH+\\\"/bboxvae-preds-\\\"+str(i)+\\\".png\\\"\\n\",\n \" )\"\n ],\n \"execution_count\": null,\n \"outputs\": []\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {\n \"id\": \"qpLk8DX9ZUoC\"\n },\n \"source\": [\n \"# Save Model and Train History\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"metadata\": {\n \"id\": \"EI4jvr7yZTrg\"\n },\n \"source\": [\n \"layoutvae.save_model(SAVE_MODEL_PATH)\\n\",\n \"layoutvae.save_history(SAVE_LOG_PATH)\"\n ],\n \"execution_count\": null,\n \"outputs\": []\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {\n \"id\": \"4wyxfkjqY3p6\"\n },\n \"source\": [\n \"# Complete Model\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"metadata\": {\n \"id\": \"zvKQB7ucLoNY\"\n },\n \"source\": [\n \"predd = layoutvae(layoutvae.test_label_set)\\n\",\n \"for i in range(2):\\n\",\n \" plot_layouts(data = predd[i*16:(i+1)*16],\\n\",\n \" colors=colors,\\n\",\n \" class_names=class_names,\\n\",\n \" path=SAVE_OUTPUT_PATH+\\\"/random-preds2-\\\"+str(i)+\\\".png\\\"\\n\",\n \" )\"\n ],\n \"execution_count\": null,\n \"outputs\": []\n }\n ]\n}" }, { "path": "LayoutVAE/Source/bboxvae.py", "content": "# -*- coding: utf-8 -*-\n\"\"\"\nCreated on Tue Aug 10 17:25:31 2021\n\n@author: Tushar & Tanishk\n\"\"\"\nimport torch as T\nfrom modelblocks import Encoder,Prior,EmbedBbox,ELBOLoss_Bbox\nfrom modelblocks import ReparamatrizeMulti,Reparamatrize_cvae,Decoder\n\nclass BboxVAE(T.nn.Module):\n def __init__(self,n_class,n_dim,max_box,latent_dim=32):\n\n super(BboxVAE,self).__init__()\n \n self.embeder = EmbedBbox(n_class)\n self.encoder = Encoder(n_dim,latent_dim=latent_dim)\n self.decoder = Decoder(n_dim,latent_dim=latent_dim)\n self.prior = Prior(latent_dim=latent_dim)\n self.loss = ELBOLoss_Bbox()\n self.rep = Reparamatrize_cvae()\n self.n_dim = n_dim\n self.n_class = n_class\n self.rep_mul = ReparamatrizeMulti()\n self.max_box = max_box\n\n\n def forward(self,inputs,isTrain=True):\n if isTrain==True :\n BoxCounts, GTBBox , BoxLabel= inputs\n los = 0\n kl1 = 0\n ll1 = 0\n for i in range(self.max_box):\n if i==0:\n PrevLabel = T.zeros((1 , *BoxLabel[... ,i,:].shape)) \n PrevBox = T.zeros((1 , *GTBBox[...,i,:].shape))\n \n\n GroundTruth = GTBBox[... , i ,:].view(-1,self.n_dim)\n \n CurrentLabel = BoxLabel[... , i ,:].view(-1,self.n_class)\n \n Embedding = self.embeder([BoxCounts,CurrentLabel,T.cat([PrevLabel,PrevBox] , dim = 2)])\n\n mu1 , logvar1 = self.encoder([GroundTruth,Embedding])\n mu2 , logvar2 = self.prior(Embedding)\n z1 = self.rep([mu1,logvar1])\n #z2 = self.rep([mu2,logvar2])\n \n Mu = self.decoder([Embedding,z1])\n BBox = self.rep_mul(Mu)\n CLoss, kl_tot , ll_tot = self.loss([mu1,logvar1,mu2,logvar2, BBox , GroundTruth])\n\n los = los + CLoss/self.max_box\n kl1 = kl1 + kl_tot/self.max_box\n ll1 = ll1 + ll_tot/self.max_box\n \n PrevBox = T.cat([PrevBox ,T.unsqueeze(GroundTruth,0)])\n PrevLabel = T.cat([PrevLabel , T.unsqueeze(CurrentLabel,0)])\n\n\n return los , kl1 , ll1\n else:\n BoxCounts, BoxLabel= inputs\n BBoxes = []\n for i in range(self.max_box):\n if i==0:\n PrevLabel = T.zeros((1 , *BoxLabel[... ,i,:].shape)) \n PrevBox = T.zeros((1 , BoxLabel.shape[0] , 4))\n\n CurrentLabel = BoxLabel[... , i ,:].view(-1,self.n_class)\n Embedding = self.embeder([BoxCounts,CurrentLabel,T.cat([PrevLabel,PrevBox] , dim = 2)])\n \n mu , logvar = self.prior(Embedding)\n \n z = self.rep([mu,logvar])\n \n Mu = self.decoder([Embedding,z])\n \n BBox = self.rep_mul(Mu)\n \n PrevBox = T.cat([PrevBox ,T.unsqueeze(BBox,0)])\n PrevLabel = T.cat([PrevLabel , T.unsqueeze(CurrentLabel,0)])\n BBoxes.append(BBox.t())\n BBoxes =T.stack(BBoxes)\n return BBoxes" }, { "path": "LayoutVAE/Source/config.py", "content": "# -*- coding: utf-8 -*-\r\n\"\"\"\r\nCreated on Tue Aug 10 18:24:51 2021\r\n\r\n@author: Tushar and Tanishk\r\n\"\"\"\r\n\r\n# PATHS\r\n\r\nSAVE_MODEL_PATH = \"\"\r\nSAVE_LOG_PATH = \"\"\r\nDATA_PATH = \"\"\r\nSAVE_OUTPUT_PATH = \"\"\r\n\r\n# Parameters\r\nCVAE_LR = 1e-5\r\nBVAE_LR = 1e-4\r\nCVAE_EPOCHS = 1\r\nBVAE_EPOCHS = 1\r\nBVAE_LATENT_DIM = 32\r\nN_CLASS = 6\r\nMAX_BOX = 9\r\nBVAE_BSIZE = 256\r\nCVAE_BSIZE = 256\r\nBVAE_VAL_SPLIT = 0.1\r\nCVAE_VAL_SPLIT = 0.1\r\nFRAC = 0.005\r\n\r\n# Other\r\nclass_names = ['None' , 'Text' , 'Title' , 'List' , 'Table' ,'Figure']\r\n" }, { "path": "LayoutVAE/Source/countvae.py", "content": "# -*- coding: utf-8 -*-\r\n\"\"\"\r\nCreated on Tue Aug 10 17:11:05 2021\r\n\r\n@author: Tushar & Tanishk\r\n\"\"\"\r\nimport torch as T\r\nfrom torch.distributions import Poisson\r\nfrom modelblocks import Encoder, Decoder, Prior, Embeder, ELBOLoss, Reparamatrize_cvae, Sampling \r\n \r\nclass CountVAE(T.nn.Module):\r\n \r\n def __init__(self,n_class,max_box=9):\r\n '''\r\n n_class = number of class (intger)\r\n max_box = maximum number of boxes (integer)\r\n isTrain(boolean) default False : defines whether data is to be treated as training data or testing\r\n \r\n if isTrain = True :\r\n input must be a tuple with first value corresponding to label set and second corresponding to ground Truth\r\n counts\r\n else :\r\n input must have label set\r\n \r\n '''\r\n super(CountVAE,self).__init__()\r\n \r\n \r\n self.encoder = Encoder()\r\n self.prior = Prior()\r\n self.decoder = Decoder(1)\r\n self.embeder = Embeder(n_class)\r\n self.loss = ELBOLoss() \r\n self.rep = Reparamatrize_cvae()\r\n self.n_class = n_class\r\n self.pois = Sampling(max_box)\r\n \r\n def forward(self, inputs, isTrain = False):\r\n \r\n if isTrain==True:\r\n \r\n label_set , groundtruth_counts = inputs\r\n Loss = 0\r\n LL = 0\r\n KL = 0\r\n previous_counts = T.zeros_like(label_set)\r\n \r\n for i in range(self.n_class):\r\n \r\n current_label = T.zeros_like(previous_counts)\r\n x_ = label_set[...,i]\r\n current_label[...,i]= x_\r\n z_ = groundtruth_counts[...,i].view(-1,1)\r\n \r\n # Generate Conditional Embedding\r\n embedding = self.embeder([label_set, current_label, previous_counts])\r\n \r\n # Encoding To latet space\r\n mu1, logvar1 = self.encoder([z_,embedding])\r\n mu2, logvar2 = self.prior(embedding)\r\n \r\n # Reparamatrized Latent variable\r\n z = self.rep([mu1,logvar1])\r\n\r\n # Decode from Latent space\r\n decoded = self.decoder([embedding,z])\r\n Closs, L_, kl_ = self.loss([mu1, logvar1, mu2, logvar2, decoded , z_])\r\n \r\n # Update Losses\r\n Loss = Loss + Closs\r\n LL = LL + L_\r\n KL = KL + kl_\r\n \r\n decoded = T.exp(decoded)\r\n \r\n # Poisson Distributions with rate of Deoded\r\n # q = self.pois(decoded)\r\n q = Poisson(decoded).sample()\r\n \r\n # update Preivious Counts\r\n previous_counts = previous_counts + current_label*(q.view(-1,1) + x_.view(-1,1))\r\n \r\n return Loss/self.n_class, KL/self.n_class, LL/self.n_class\r\n \r\n else:\r\n \r\n label_set = inputs\r\n previous_counts = T.zeros_like(label_set)\r\n \r\n for i in range(self.n_class):\r\n\r\n current_label = T.zeros_like(previous_counts)\r\n x_ = label_set[...,i]\r\n current_label[...,i]= x_\r\n \r\n \r\n # Generate Conditional Embedding\r\n embedding = self.embeder([label_set, current_label, previous_counts])\r\n \r\n # Encoding To latet space\r\n mu,logvar = self.prior(embedding)\r\n \r\n # Reparamatrized Latent variable\r\n z = self.rep([mu,logvar])\r\n \r\n # Decode from Latent space\r\n decoded = self.decoder([embedding,z])\r\n decoded = T.exp(decoded)\r\n\r\n # Poisson Distributions with rate of Deoded\r\n # q = self.pois(decoded)\r\n q = Poisson(decoded).sample()\r\n \r\n # update Preivious Counts\r\n previous_counts = previous_counts + current_label*(q.view(-1,1) + x_.view(-1,1))\r\n \r\n return previous_counts" }, { "path": "LayoutVAE/Source/layoutvae.py", "content": "# -*- coding: utf-8 -*-\n\"\"\"\nCreated on Tue Aug 10 17:44:46 2021\n\n@author: Tushar & Tanishk\n\"\"\"\nimport torch as T\nfrom countvae import CountVAE\nfrom bboxvae import BboxVAE\nimport matplotlib.pyplot as plt\nimport pandas as pd\nimport numpy as np\n\n########################\n###### LAYOUT VAE ######\n########################\n\n\nclass LayoutVAE(T.nn.Module):\n\n def __init__(self, n_class = 6, max_box = 9,bboxvae_latent_dim = 32,bboxvae_lr=1e-4,countvae_lr=1e-6):\n '''\n ** Layout VAE **\n * https://arxiv.org/abs/1907.10719\n '''\n super(LayoutVAE,self).__init__()\n\n self.max_box = max_box\n self.n_class = n_class\n self.lr_bvae = bboxvae_lr\n self.lr_cvae = countvae_lr\n self.countvae = CountVAE(n_class)\n self.bboxvae = BboxVAE(n_class,4,max_box,bboxvae_latent_dim)\n self.is_cvae_trained = 0\n self.is_bvae_trained = 0\n\n def forward(self,input):\n '''\n Takes only Labels Set as input\n Label Set : it is a vector of size n_class and contains 1 if correspinding class is present\n '''\n if self.is_cvae_trained == 0:\n print(\"[Warning] Count VAE is Not Trained !!\")\n\n if self.is_bvae_trained == 0:\n print(\"[Warning] Bbox VAE is Not Trained !!\")\n\n label_set = input\n pred_class_counts = self.countvae(label_set , isTrain=False)\n\n # Normalize classiction between [0 , max_box]\n pred_class_counts = T.floor ( self.max_box*(pred_class_counts / T.sum(pred_class_counts , dim = 1 ).view(-1,1)) )\n\n # Extra boxes which are not be predicted\n # Their counts are set in first class\n for class_count in pred_class_counts:\n if(T.sum(class_count) < self.max_box):\n class_count[0] = self.max_box - T.sum(class_count)\n\n class_labels = T.zeros(len(label_set) , self.max_box, self.n_class)\n\n for i in range(len(pred_class_counts)):\n l = 0\n for j in range(self.n_class):\n for k in range(int(pred_class_counts[i][self.n_class-j-1])):\n class_labels[i][l][self.n_class-j-1] = 1;\n l+=1\n\n pred_box = self.bboxvae([ pred_class_counts, class_labels], isTrain=False)\n pred_box = pred_box.permute(2,0,1)\n class_info = T.unsqueeze(T.argmax(class_labels ,dim=2),dim=2)\n predictions = T.cat([class_info,pred_box],dim = 2)\n\n for i in range(len(predictions)):\n for j in range(len(predictions[i])):\n if predictions[i][j][0]==0:\n predictions[i][j]*=0\n\n self.predictions = predictions\n self.pred_class_counts = pred_class_counts\n\n return predictions\n\n def load_data(self, path, frac = 0.5, train_test_split = 0.1):\n '''\n Loads data from npy file\n path string containig path to data\n frac defines the fraction of data to load\n\n '''\n try :\n Data = np.load(path)\n # Sortind Data in proper order\n np.random.shuffle(Data)\n order = np.argsort(Data[:,:,0])\n for i in range(len(Data)):\n Data[i] = Data[i][order[i][::-1]]\n\n data_size = int(frac*len(Data))\n test_size = int(train_test_split*data_size)\n Data = T.tensor(Data[0:data_size]).float()\n test_data = Data[0:test_size]\n Data = Data[test_size:]\n\n # Prepare Data\n self.class_labels = Data[...,4:]\n self.class_counts = T.sum(Data[...,4:], dim = 1)\n self.b_boxes = Data[...,0:4]\n self.label_set = (self.class_counts !=0) + 0.0\n\n # Test Data\n self.test_class_labels = test_data[...,4:]\n self.test_class_counts = T.sum(test_data[...,4:], dim = 1)\n self.test_b_boxes = test_data[...,0:4]\n self.test_label_set = (self.test_class_counts !=0) + 0.0\n\n print(\"[Success] Data Loaded Succesfully\")\n\n except:\n print(\"[Failed] Data Loading Failed\\n please check path\")\n\n def train(self, optim, train_mode = 'bboxvae', epochs = 100, bsize = 256 , validation_split = 0.1):\n '''\n * train_mode (str , default bboxvae) : Two optons\n 1. if train_mode is bboxvae, BBoxVAE model will be trained and data\n will be loaded accordingly\n 2. if train_mode is countvae, CountVAE model will be trained and data\n will be loaded accordingly\n * epochs (int , default 100 ) : number of epochs training should run\n * bsize(int default 256) : Batch Size\n * validation_split(float default 0.1) : should be between between 0 and 1\n 1 . it defines the size of validation data\n\n '''\n # Create validation Split\n total_examples = len(self.class_counts)\n val_size = int(total_examples*validation_split)\n\n losses = dict()\n train_data = []\n if train_mode == 'countvae':\n model = self.countvae\n train_data = [self.label_set, self.class_counts]\n else :\n model = self.bboxvae\n train_data = [self.class_counts, self.b_boxes, self.class_labels]\n\n # Validation Data\n val_data = []\n for x in train_data:\n val_data.append(x[:val_size])\n\n # Train data\n for i in range(len(train_data)):\n train_data[i] = train_data[i][val_size:]\n\n\n # find the number of batches\n batches = len(train_data[0])//bsize\n second_loss = 'mse'\n if train_mode == 'countvae':\n second_loss = 'poisson_nll'\n\n # Dictionary to keep track of model statistics\n losses = {'epoch':-1,\n 'batch':0,\n 'lr' : 0,\n 'loss':0,\n 'kl_div_loss':0,\n second_loss+'_loss':0,\n 'val_loss':0,\n 'val_kl_div_loss':0,\n 'val_'+second_loss+'_loss':0\n }\n\n history = pd.DataFrame(losses ,index = [0])\n index = 1\n\n for ep in range(epochs):\n\n # if train_mode=='countvae':\n # self.countvae_pred_grpah(epoch = ep,path = CVAE_PATH)\n\n print(f'Epoch[{ep+1}/{epochs}]')\n for batch in range(batches):\n\n # Get Current batch\n b = []\n for x in train_data:\n b.append(x[batch*bsize : (batch+1)*bsize])\n\n optim.zero_grad()\n\n # Train Step\n loss, kl_, l_ = model(b,isTrain = True)\n\n # Validation Step\n val_loss, val_kl_, val_l_ = model(val_data, isTrain = True)\n\n\n # Save Statistics\n losses['epoch'] = ep\n losses['batch'] = batch\n losses['lr'] = optim.param_groups[0]['lr']\n\n loss_list = [loss, kl_, l_ , val_loss , val_kl_ , val_l_]\n\n for i in range(6):\n losses[list(losses.keys())[3+i]] = loss_list[i].cpu().clone().detach().numpy()\n pass\n\n losses_df = pd.DataFrame(losses , index=[index])\n history = pd.concat([history,losses_df])\n index+=1\n\n # Backpropogation step and updating weights\n loss.backward()\n optim.step()\n print('\\r Batch: {}/{} - loss : {} - val_loss : {} - val_{} : {}'.format(batch+1,batches,\n losses_df['loss'][index-1],\n losses_df['val_loss'][index-1],\n second_loss,\n losses_df['val_'+second_loss+'_loss'][index-1]),\n end=\"\")\n print(\"\\n\")\n print('[Success] Finished Training')\n return history\n\n def load_countvae_weights(self,path):\n try :\n self.countvae = T.load(path)\n self.is_cvae_trained=1\n print('[Success] Loaded Successfully')\n except:\n print('[Failed] Load Failed')\n\n def load_bboxvae_weights(self,path):\n try :\n self.bboxvae = T.load(path)\n self.is_bvae_trained=1\n print('[Success] Loaded Successfully')\n except:\n print('[Failed] Load Failed')\n\n def train_bboxvae(self,epochs=30, bsize=256, validation_split=0.1, optim=None):\n if optim == None:\n optim = T.optim.Adam(self.bboxvae.parameters(),lr=self.lr_bvae)\n\n # Start Training\n history = self.train(optim = optim,\n train_mode = 'bboxvae',\n epochs = epochs,\n bsize = bsize,\n validation_split = validation_split\n )\n self.is_bvae_trained = 1\n self.bvae_history = history[history.columns][1:]\n return self.bvae_history\n\n def train_countvae(self,epochs=30, bsize=256, validation_split=0.1, optim=None):\n\n if optim == None:\n optim = T.optim.Adam(self.countvae.parameters(),lr=self.lr_cvae)\n\n # Start Training\n history = self.train(optim = optim,\n train_mode = 'countvae',\n epochs = epochs,\n bsize = bsize,\n validation_split = validation_split\n )\n self.is_cvae_trained = 1\n self.cvae_history = history[history.columns][1:]\n return self.cvae_history\n\n def pred_countvae(self,data=None):\n '''\n * Functions is used for for predcting from CountVAE\n given label_set\n * if data is None than label set from loaded data\n are used for predictions.\n '''\n\n if self.is_cvae_trained == 0:\n print(\"[Warning] Count VAE is Not Trained !!\")\n if data == None :\n data = self.test_label_set\n return self.countvae(data , isTrain=False)\n\n def pred_bboxvae(self,Data=None):\n\n '''\n * Functions is used for for predcting from BboxVAE\n given class_counts and class labels\n * if data is None than class counts and class labels from loaded data\n are used for predictions.\n '''\n\n if self.is_bvae_trained == 0:\n print(\"[Warning] Bbox VAE is Not Trained !!\")\n\n if Data == None :\n Data = [self.test_class_counts,self.test_class_labels]\n\n batches = len(Data[0])//64\n\n for b in range(batches):\n\n # Get data in batch\n data = [self.test_class_counts[b*64 : (b+1)*64],\n self.test_class_labels[b*64 : (b+1)*64]]\n\n # Predict\n pred = self.bboxvae(data, isTrain=False)\n pred = pred.permute(2,0,1)\n\n # cxywh format\n class_info = T.unsqueeze(T.argmax(data[1] ,dim=2),dim=2)\n pred = T.cat([class_info,pred],dim = 2)\n\n\n for i in range(len(pred)):\n for j in range(len(pred[i])):\n if pred[i][j][0]==0:\n pred[i][j] *= 0\n\n if b > 0:\n predictions = T.cat([predictions,pred],dim=0)\n else:\n predictions = pred\n class_info =T.argmax(self.test_class_labels[0:64*batches] ,dim=2)\n class_info = T.unsqueeze(class_info,dim=2)\n gt = T.cat([class_info,self.test_b_boxes[0:64*batches]],dim = 2)\n return predictions, gt\n\n def countvae_pred_grpah(self,path,epoch = 0):\n pred_cvae = self.pred_countvae()\n pred_cvae = T.sum(pred_cvae,dim=0)\n pred_cvae = pred_cvae/T.sum(pred_cvae)\n pred_cvae = pred_cvae.clone().detach().numpy()\n\n gt_cvae = T.sum(self.class_counts,dim=0)\n gt_cvae = gt_cvae/T.sum(gt_cvae)\n gt_cvae = gt_cvae.clone().detach().numpy()\n\n fig = plt.figure(figsize=(5 ,4), dpi=100 ,facecolor=(0,0,0))\n ax = fig.add_subplot()\n ax.plot(gt_cvae , 'red',marker = 'o', label = 'Ground Truth',linewidth=4)\n ax.plot(pred_cvae,'blue',marker ='o',label = \"Predicted\" ,linewidth=4)\n ax.legend()\n ax.set_title('Ground Truth vs Predicted Distribution\\n Epoch = '+str(epoch))\n ax.set_xlabel('Classes')\n ax.set_xticks([i for i in range(config.N_CLASS)])\n ax.set_xticklabels(config.class_names)\n\n plt.savefig(path+\"cvae-train-ep-\"+str(epoch)+\".png\",facecolor=(0,0,0))\n plt.close()\n\n\n def convert_to_cxywh(self,data):\n '''\n\n\n Parameters\n ----------\n data : (torch.tensor) tensor\n tensor of size (N , B , 4 + C)\n N = number of examples\n B = Number of boxes\n C = Number of classes\n\n Returns\n -------\n cxywh : (torch.tensor) tensor\n tensor of size (N , B , 1 + 4)\n N = number of examples\n B = Number of boxes\n c = class\n (x,y) = upper left corner\n w and h = height and width\n\n '''\n bboxes = data[...,0:4]\n labels = data[...,4: ]\n class_info = T.unsqueeze(T.argmax(labels ,dim=2),dim=2)\n cxywh = T.cat([class_info,bboxes],dim = 2)\n return cxywh\n\n def save_model(self,path):\n\n T.save(self.countvae,path+'countvae.h5')\n T.save(self.bboxvae,path+'bboxvae.h5')\n T.save(self,path+'layoutvae.h5')\n print('[Success] Saved Successfully')\n\n def save_history(self,path):\n\n self.cvae_history.to_csv(path+'cvae-history.csv',index=False)\n self.bvae_history.to_csv(path+'bvae-history.csv',index=False)\n print('[Success] Saved Successfully')\n" }, { "path": "LayoutVAE/Source/main.py", "content": "# -*- coding: utf-8 -*-\r\n\"\"\"\r\nCreated on Tue Aug 10 18:26:19 2021\r\n\r\n@author: Tushar & Tanishk\r\n\"\"\"\r\nfrom layoutvae import LayoutVAE\r\nfrom utils import plot_layouts ,plot_history,generate_colors,countvae_pred_graph\r\nimport config\r\n\r\n\r\n# Model\r\nlayoutvae = LayoutVAE(n_class=config.N_CLASS, \r\n max_box=config.MAX_BOX,\r\n bboxvae_latent_dim=config.BVAE_LATENT_DIM,\r\n bboxvae_lr=config.BVAE_LR,\r\n countvae_lr=config.CVAE_LR,\r\n )\r\nlayoutvae.load_data(path = config.DATA_PATH,frac= config.FRAC)\r\n\r\nhistory_bvae_df = layoutvae.train_bboxvae(bsize = config.BVAE_BSIZE,\r\n epochs=config.BVAE_EPOCHS,\r\n validation_split=config.BVAE_VAL_SPLIT)\r\n\r\n\r\nhistory_cvae_df = layoutvae.train_countvae(bsize = config.CVAE_BSIZE,\r\n epochs=config.CVAE_EPOCHS,\r\n validation_split=config.CVAE_VAL_SPLIT)\r\n\r\n# Save History\r\nlayoutvae.save_model(config.SAVE_MODEL_PATH)\r\nlayoutvae.save_history(config.SAVE_LOG_PATH)\r\n\r\n# Predict Layout\r\ncolors = generate_colors(n_class=config.N_CLASS,\r\n class_names=config.class_names)\r\n\r\n# only using bboxvae\r\npred , ground_truth = layoutvae.pred_bboxvae()\r\nfor i in range(2):\r\n plot_layouts(pred = pred[i*16:(i+1)*16],\r\n colors=colors,\r\n class_names=config.class_names,\r\n path=config.SAVE_OUTPUT_PATH+\"bvae-preds-\"+str(i)+\".png\"\r\n )\r\n \r\n# using complete model\r\nfinal_predictions = layoutvae(layoutvae.label_set)\r\n\r\n#visualize and save predictions\r\n\r\n\r\nplot_layouts(pred = final_predictions,\r\n colors = colors,\r\n title = \"Random Outputs\",\r\n class_names=config.class_names,\r\n path = config.SAVE_OUTPUT_PATH+\"randout.svg\")\r\n\r\ncountvae_pred_graph(layoutvae,config.SAVE_OUTPUT_PATH+\"cvae-train.png\")\r\n\r\nplot_layouts(pred = pred,\r\n colors = colors,\r\n title = \"BBoxVAE Outputs\",\r\n class_names=config.class_names,\r\n path = config.SAVE_OUTPUT_PATH+\"bboxvae-pred.svg\")\r\n\r\n# Plot and save Train History plots\r\nplot_history(layoutvae.bvae_history , path = config.SAVE_LOG_PATH+\"bvae-train.png\")\r\nplot_history(layoutvae.cvae_history , path = config.SAVE_LOG_PATH+\"cvae-train.png\")\r\n\r\n# Complete Model\r\npredd = layoutvae(layoutvae.test_label_set)\r\nfor i in range(2):\r\n plot_layouts(pred = predd[i*16:(i+1)*16],\r\n colors=colors,\r\n class_names=config.class_names,\r\n path=config.SAVE_OUTPUT_PATH+\"random-preds-\"+str(i)+\".png\"\r\n )\r\n" }, { "path": "LayoutVAE/Source/modelblocks.py", "content": "# -*- coding: utf-8 -*-\n\"\"\"\nCreated on Tue Aug 10 17:14:02 2021\n\n@author: Tushar & Tanishk\n\"\"\"\nfrom __future__ import division\nimport torch as T\nfrom torch.nn import Sequential , Linear , ReLU , PoissonNLLLoss, LSTM\nfrom torch.distributions import MultivariateNormal \n\nclass fcblock(T.nn.Module):\n def __init__(self, n_class):\n super(fcblock, self).__init__()\n self.seq = Sequential(\n Linear(n_class,128),\n ReLU(),\n Linear(128,128),\n ReLU(),\n )\n def forward(self,inputs):\n out = self.seq(inputs)\n return out\n\nclass Embeder(T.nn.Module):\n def __init__(self,n_class):\n super(Embeder,self).__init__()\n \n self.fcb1 = fcblock(n_class)\n self.fcb2 = fcblock(n_class)\n self.fcb3 = fcblock(n_class)\n self.fc = Linear(128*3,128)\n\n \n def forward(self,inputs):\n in1,in2,in3 = inputs\n in1 = self.fcb1(in1)\n in2 = self.fcb2(in2)\n in3 = self.fcb3(in3)\n out = T.cat((in1,in2,in3),1)\n out = self.fc(out)\n return out\n\nclass Encoder(T.nn.Module):\n def __init__(self, in_dim=1 ,latent_dim=32):\n super(Encoder,self).__init__()\n self.act = ReLU()\n self.fc1 = Linear(in_dim,128)\n self.fc2 = Linear(128,128)\n self.fc3 = Linear(256,latent_dim)\n self.fc4 = Linear(latent_dim,latent_dim)\n self.fc5 = Linear(latent_dim,latent_dim)\n \n def forward(self,inputs):\n in1,in2 = inputs\n out = self.fc1(in1)\n out = self.act(out)\n out = self.fc2(out)\n out = T.cat((out,in2),1)\n out = self.fc3(out)\n out = self.act(out)\n mu = self.fc4(out)\n logvar = self.fc5(out)\n return mu,logvar\n\n\nclass Prior(T.nn.Module):\n def __init__(self,latent_dim=32):\n super(Prior,self).__init__()\n \n self.act = ReLU()\n self.fc1 = Linear(128,latent_dim)\n self.fc2 = Linear(latent_dim,latent_dim)\n self.fc3 = Linear(latent_dim,latent_dim)\n \n def forward(self,inputs):\n out = inputs\n out = self.fc1(out)\n out = self.act(out)\n mu = self.fc2(out)\n logvar = self.fc3(out) \n return mu,logvar\n\nclass Decoder(T.nn.Module):\n def __init__(self,output_dim,latent_dim=32):\n super(Decoder,self).__init__()\n self.act = ReLU()\n self.fc1 = Linear(128+latent_dim,128)\n self.fc2 = Linear(128,64)\n self.fc3 = Linear(64,output_dim)\n \n def forward(self,inputs):\n in1,in2 = inputs\n out = T.cat((in1,in2),1)\n out = self.fc1(out)\n out = self.act(out)\n out = self.fc2(out)\n out = self.act(out)\n out = self.fc3(out)\n return out\n\n\"\"\"# LOSS FUNCTION\"\"\"\n\nclass ELBOLoss(T.nn.Module):\n\n def __init__(self):\n super(ELBOLoss,self).__init__()\n \n def forward(self,inputs):\n mu1, logvar1, mu2, logvar2 , in1, in2 = inputs\n\n mask = (in2>0)+0.0\n in2 = in2-mask\n\n '''KL Divergence'''\n kl = 0.5 * T.sum((logvar2 - logvar1) - 1 + (logvar1.exp() + (mu2 - mu1).pow(2) )/logvar2.exp() , dim = 1).mean()\n \n '''Poisson Negative Log Likelihood'''\n pnll = PoissonNLLLoss()(in1,in2)\n\n loss = kl+pnll\n \n return loss, pnll , kl\n \n\n\nclass EmbedBbox(T.nn.Module):\n \n def __init__(self,n_class):\n super(EmbedBbox,self).__init__()\n \n self.fcb1 = fcblock(n_class)\n self.fcb2 = fcblock(n_class)\n self.seq1 = Sequential(\n Linear(128,128),\n ReLU()\n )\n \n self.n_class = n_class\n self.fc = Linear(128*3,128)\n self.lstm = LSTM(n_class+4, hidden_size=128)\n\n def forward(self,inputs):\n \n in1,in2,in3 = inputs\n\n _ , (h_0 , c_0 ) = self.lstm(in3)\n hn = h_0.view(-1, 128)\n \n in1 = self.fcb1(in1)\n in2 = self.fcb2(in2)\n in3 = self.seq1(hn)\n \n out = T.cat((in1,in2,in3),1)\n out = self.fc(out)\n \n return out\n \nclass ELBOLoss_Bbox(T.nn.Module):\n \n def __init__(self):\n super(ELBOLoss_Bbox,self).__init__()\n \n def forward(self,inputs):\n mu1,logvar1,mu2,logvar2, xp , yp = inputs\n \n ''' KL Divergence '''\n kl = 0.5 * T.sum((logvar2 - logvar1) - 1 + (logvar1.exp() + (mu2 - mu1).pow(2) )/logvar2.exp() , dim = -1 ).mean()\n \n ''' Multivariate Guassian Likelihood '''\n mse = T.nn.MSELoss()(xp,yp)\n loss = mse + kl\n \n return loss, kl,mse\n\n\nclass Reparamatrize_bvae(T.nn.Module):\n \n def __init__(self):\n super(Reparamatrize_bvae,self).__init__()\n \n def forward(self,inputs):\n \n mu , logvar = inputs\n std = T.exp(logvar/2)\n eps = T.rand_like(std)\n\n return eps*std + mu\n \n\nclass ReparamatrizeMulti(T.nn.Module):\n \n def __init__(self):\n super(ReparamatrizeMulti,self).__init__()\n \n def forward(self,inputs):\n \n mu = inputs\n std = (T.ones_like(mu)*0.02)\n eps = T.rand_like(std)\n \n return eps*std + mu\n \nclass Reparamatrize_cvae(T.nn.Module):\n \n def __init__(self):\n super(Reparamatrize_cvae,self).__init__()\n \n def forward(self,inputs):\n \n mu , logvar = inputs\n '''\n mu = mean \n logvar = log of diagonal elements of covariance matrix\n '''\n # Covarince Matrix\n covar = T.diag_embed(T.exp(logvar/2), dim1=-2,dim2=-1)\n\n # Multivariate Normal Distribution\n p = MultivariateNormal(mu,covar)\n z_latent = p.rsample().float()\n return z_latent\n\nclass Sampling(T.nn.Module):\n\n def __init__(self,MAX_BOX):\n super(Sampling,self).__init__()\n self.max_box = MAX_BOX\n \n def forward(self,lamda):\n \n lamda = lamda.view(-1)\n mask = T.zeros(lamda.shape[0] , self.max_box)\n lamda = T.t(T.t(mask) + lamda)\n mask = mask + T.arange(0,self.max_box,1)\n e_lamda = T.exp(lamda)\n lamda_x = lamda ** mask \n fact = T.exp(T.lgamma(T.arange(0 , self.max_box)+1))\n \n # P = ((lambda ^ x)*e^(lamda)) / x! \n probab = (lamda_x*e_lamda)/fact\n sample = T.argmax(probab,dim=1)\n\n return sample" }, { "path": "LayoutVAE/Source/utils.py", "content": "# -*- coding: utf-8 -*-\n\"\"\"\nCreated on Tue Aug 10 17:36:04 2021\n\n@author: Tushar & Tanishk\n\"\"\"\n\nimport torch as T\nimport matplotlib.pyplot as plt\nimport matplotlib.patches as patches\nfrom matplotlib.patches import Patch\nimport config\nplt.style.use('dark_background')\n\ndef plot_history(history,title = 'Training Statistics', path =\"\"):\n height = 12\n width = 9\n fig = plt.figure(figsize=(width,height), dpi=100 ,facecolor=(0,0,0))\n height_ratio = [0.25,1,1,1]\n grid = plt.GridSpec(4,2,\n hspace=0.3,wspace=0.2,\n height_ratios =height_ratio,\n left=0.02,right=0.98,top=0.98,bottom=0.02\n )\n index = 0\n ax = fig.add_subplot(grid[index : index+2])\n index+=2\n ax.text(x = 0.3 ,y = 0.5 ,s = title,fontsize=30)\n ax.invert_yaxis()\n ax.axis('off')\n colors = ['red','blue','green']\n for i in range(3):\n\n ax = fig.add_subplot(grid[index])\n ax.plot(history[history.columns[i+3]],colors[i])\n index+=1\n ax.set_facecolor((0,0,0))\n ax.set_title(history.columns[i+3])\n ax = fig.add_subplot(grid[index])\n ax.plot(history[history.columns[i+6]],colors[i])\n ax.set_title(history.columns[i+6])\n index+=1\n ax.set_facecolor((0,0,0))\n ax.set_xlabel('Batches')\n ax.set_ylabel('Loss')\n plt.savefig(path, facecolor=(0,0,0))\n\n\ndef generate_colors(class_names = None,n_class=6):\n '''\n\n Parameters\n ----------\n class_names : list, optional\n List of classes in the dataset. The default is None.\n n_class : integer, optional\n The default is 6.\n Returns\n -------\n colors : list of hexadecimal strings\n \n '''\n cmap = [\"\",\"#dc143c\",\"#ffff00\",\"#00ff00\",\"#ff00ff\",\"#1e90ff\",\"#fff5ee\",\n \"#00ffff\",\"#8b008b\",\"#ff4500\",\"#8b4513\",\"#808000\",\"#483d8b\",\n \"#008000\",\"#000080\",\"#9acd32\",\"#ffa500\",\"#ba55d3\",\"#00fa9a\",\n \"#dc143c\",\"#0000ff\",\"#f08080\",\"#f0e68c\",\"#dda0dd\",\"#ff1493\"]\n \n colors = dict()\n\n if class_names == None:\n class_names = []\n for i in range(n_class):\n class_names.append('class'+str(i+1))\n \n for i in range(n_class):\n colors[class_names[i]] = cmap[i]\n\n return colors\n\nclass_names = ['None' , 'Text' , 'Title' , 'List' , 'Table' ,'Figure']\ncolors = generate_colors(n_class=6 , class_names=class_names)\n\ndef plot_layouts(pred,colors,class_names,title=\"Predictions\", path=\"\"):\n '''\n data in cxywh format\n '''\n height = 15\n width = 9\n fig = plt.figure(figsize=(width,height), dpi=50 ,facecolor=(0,0,0))\n height_ratio = [0.25,0.25,1,1,1,1]\n grid = plt.GridSpec(6,4,\n hspace=0.05,wspace=0.05,\n height_ratios =height_ratio,\n left=0.02,right=0.98,top=0.98,bottom=0.02\n )\n index = 0\n\n\n ax = fig.add_subplot(grid[index : index+4])\n index+=4\n ax.text(x = 0.2 ,y = 0.5 ,s = title,fontsize=30)\n\n legend = []\n ax = fig.add_subplot(grid[index : index+4])\n index += 4\n \n for i in range(1,6):\n legend.append(Patch(facecolor=colors[class_names[i]]+\"40\",\n edgecolor=colors[class_names[i]],\n label= class_names[i]))\n \n ax.legend(handles=legend, ncol=3,loc=8, fontsize=25, facecolor=(0,0,0))\n ax.axis('off')\n\n for i in range(16):\n ax = fig.add_subplot(grid[index])\n index += 1\n \n data = pred[i]\n rect1 = patches.Rectangle((0,0),180,240)\n rect1.set_color((0,0,0,1))\n ax.add_patch(rect1)\n for box in data:\n\n c,x,y,w,h = box\n if c==0:\n continue\n x = x*180\n y = y*240\n w = w*180\n h = h*240\n rect = patches.Rectangle((x,y),w,h,linewidth=2)\n rect.set_color(colors[class_names[int(c)]]+\"72\")\n rect.set_linestyle('-')\n rect.set_edgecolor(colors[class_names[int(c)]])\n ax.add_patch(rect)\n ax.plot()\n ax.set_facecolor((0,0,0))\n for spine in ax.spines.values():\n spine.set_edgecolor('green')\n spine.set_linewidth(2)\n ax.invert_yaxis()\n ax.set_xticks([])\n ax.set_yticks([])\n plt.savefig(path , facecolor=(0,0,0))\n\ndef countvae_pred_graph(model,path=\"\"):\n pred_cvae = model.pred_countvae()\n pred_cvae = T.sum(pred_cvae,dim=0)\n pred_cvae = pred_cvae/T.sum(pred_cvae)\n pred_cvae = pred_cvae.to('cpu').clone().detach().numpy()\n\n gt_cvae = T.sum(model.class_counts,dim=0)\n gt_cvae = gt_cvae/T.sum(gt_cvae)\n gt_cvae = gt_cvae.to('cpu').clone().detach().numpy()\n\n fig = plt.figure(figsize=(5 ,4), dpi=100 ,facecolor=(0,0,0))\n ax = fig.add_subplot()\n ax.plot(gt_cvae , 'red',marker = 'o', label = 'Ground Truth',linewidth=4)\n ax.plot(pred_cvae,'blue',marker ='o',label = \"Predicted\" ,linewidth=4)\n ax.legend()\n ax.set_title('Ground Truth vs Predicted Distribution')\n ax.set_xlabel('Classes')\n ax.set_xticks([i for i in range(config.N_CLASS)])\n ax.set_xticklabels(config.class_names)\n\n plt.savefig(path,facecolor=(0,0,0))\n plt.close()" }, { "path": "LayoutVAE/readme.md", "content": "# Layout VAE\n\n## Introduction\nThis repository provides PyTorch 1.9.0 implementation of Layout VAE [[1]](#1) ( Layout Variational Auto Encoder ). It is a probabilistic and autoregressive model which generates the scene layout using latent variables in lower dimensions . It is capable of generating different layouts using the same data point.\n\n## Architecture\n![Architecture](https://user-images.githubusercontent.com/40228110/129759853-cc9595bc-9a96-4e68-ac7b-97112cdae528.png)\n\n\n## Requirements\n- PyTorch 1.9.0\n- Python 3.8\n\n## Datasets\n* **PubLayNet** : It is a dataset for document layout analysis. It contains images of research papers and articles and annotations for various elements in a page such as “text”, “list”, “figure” etc in these research paper images. The dataset was obtained by automatically matching the XML representations and the content of over 1 million PDF articles that are publicly available on PubMed Central.[[2]](#2)\n* We sorted the cordinates of Bounding boxes in left to right direction. and used 50% data as train data and 5% test data. \n\n## Getting Started\n\"Open\n\nClick on the above Badge to quickstart the LayoutVAE in google colab\n\n## Results\n### CountVAE\n![](Outputs/countvae/countvae-training.gif)\n\n### BBoxVAE\n\"Image\"\n\n## References\n[1]\nLayoutVAE: Stochastic Scene Layout Generation from a Label Set.Akash Abdu Jyothi and Thibaut Durand and Jiawei He and Leonid Sigal and Greg Mori [Paper](https://arxiv.org/abs/1907.10719)\n\n[2]\nPubLayNet : By Xu Zhong, Jianbin Tang, Antonio Jimeno Yepes [Dataset](https://developer.ibm.com/exchanges/data/all/publaynet/)\n" }, { "path": "Metrics/README.md", "content": "# Metrics/Quantitative Comparison\r\n## Intersection over Union (IoU)\r\nThe intersection over the union of boxes is calculated pairwise and are then added together. The overall IoU of the data is averaged over all the documents.\r\n\r\nFor the kth document in the data, the iou Lk is calculated as follows:\r\n\r\n![iou1](/Metrics/readme_images/iou1.jpg)\r\n\r\nWhere n is the total number of boxes in the document.\r\n\r\nFor the whole data, the loss(IoU) is calculated as follows:\r\n\r\n![iou2](/Metrics/readme_images/iou2.jpg)\r\n\r\nWhere N is the total number of documents in the data.\r\n\r\n\r\n## Overlapping Loss\r\nOverlapping loss is defined as the ratio of overlapping area by the box area. It is also calculated pairwise, added together and then averaged for all documents. Related expressions are given below:\r\n\r\n![overlapping1](/Metrics/readme_images/overlapping1.jpg)\r\n\r\n![overlapping2](/Metrics/readme_images/iou2.jpg)\r\n\r\n## Alignment Loss\r\nAdjacent elements (boxes) are usually in six possible alignment types: Left, X-center, Right, Top, Y-center and Bottom aligned. Denote =(xL,yT,xC,yC,xR,yB) as the top-left, center and bottom-right coordinates of the predicted bounding box, we encourage pairwise alignment among elements by introducing an alignment loss:\r\n\r\n![alg1](/Metrics/readme_images/algn1.jpg)\r\n\r\n![alg2](/Metrics/readme_images/algn2.jpg)\r\n\r\n## Comparison\r\nData was normalised with respect to the original data.\r\n| | Overlap | IOU | Alignment |\r\n|--------------------|:-----------:|:-----------:|:---------:|\r\n| Original Data | 1.000000 | 1.000000 | 1.000000 |\r\n| LayoutGAN | 1172.005234 | 2745.437529 | 1.164882 |\r\n| LayoutVAE | 119.320127 | 185.864381 | 3.493406 |\r\n| Layout Transformer | 1.090315 | 1.422297 | 0.739862 |\r\n" }, { "path": "Metrics/metrics.ipynb", "content": "{\n \"nbformat\": 4,\n \"nbformat_minor\": 2,\n \"metadata\": {\n \"colab\": {\n \"name\": \"metrics.ipynb\",\n \"provenance\": [],\n \"collapsed_sections\": []\n },\n \"kernelspec\": {\n \"name\": \"python3\",\n \"display_name\": \"Python 3\"\n },\n \"language_info\": {\n \"name\": \"python\"\n }\n },\n \"cells\": [\n {\n \"cell_type\": \"markdown\",\n \"source\": [\n \"#Imports\"\n ],\n \"metadata\": {\n \"id\": \"F_A3n24sMYfx\"\n }\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 1,\n \"source\": [\n \"# Drive Mounting (for Google Colab only). If not using Colab, comment the below two lines.\\r\\n\",\n \"from google.colab import drive\\r\\n\",\n \"drive.mount('/content/drive')\"\n ],\n \"outputs\": [\n {\n \"output_type\": \"stream\",\n \"name\": \"stdout\",\n \"text\": [\n \"Drive already mounted at /content/drive; to attempt to forcibly remount, call drive.mount(\\\"/content/drive\\\", force_remount=True).\\n\"\n ]\n }\n ],\n \"metadata\": {\n \"colab\": {\n \"base_uri\": \"https://localhost:8080/\"\n },\n \"id\": \"I3oR3QRKPSfu\",\n \"outputId\": \"e9e4bcd8-0436-48d4-ed38-a338a4ae3d12\"\n }\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 2,\n \"source\": [\n \"import numpy as np\\r\\n\",\n \"import pandas as pd\\r\\n\",\n \"from shapely.geometry import Polygon\"\n ],\n \"outputs\": [],\n \"metadata\": {\n \"id\": \"yEjKDwojt5qE\"\n }\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 3,\n \"source\": [\n \"root=\\\"/content/drive/MyDrive/Folder_Name/\\\" #path of the root directory\"\n ],\n \"outputs\": [],\n \"metadata\": {\n \"id\": \"X3kiwgIpsQzN\"\n }\n },\n {\n \"cell_type\": \"markdown\",\n \"source\": [\n \"#Data Loading\"\n ],\n \"metadata\": {\n \"id\": \"jdCPjnwvMgVc\"\n }\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 4,\n \"source\": [\n \"publaydata = np.load(root+\\\"publaynet.npy\\\")\\r\\n\",\n \"Transformer_res=np.load(root+\\\"trans.npy\\\")\\r\\n\",\n \"VAE_res = np.load(root+\\\"VAE_res.npy\\\")\\r\\n\",\n \"GAN_res = np.load(root+\\\"GAN_res.npy\\\")\"\n ],\n \"outputs\": [],\n \"metadata\": {\n \"id\": \"3ybVHIU8uRoW\"\n }\n },\n {\n \"cell_type\": \"markdown\",\n \"source\": [\n \"Reshaping and arranging data in an optimal format. Preferred format is [c x y w h]\\n\",\n \"\\n\",\n \"* c is the class of the box.\\n\",\n \"* x and y are the corrdinates for the top left corner of the box.\\n\",\n \"* w and h are the width and height respectively. \\n\",\n \"---\\n\",\n \"\\n\"\n ],\n \"metadata\": {\n \"id\": \"Y83pq-l_NL8y\"\n }\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 5,\n \"source\": [\n \"g_data = GAN_res.reshape((1024, 9, 9))[:1000]\\r\\n\",\n \"g_data = g_data[...,0:5]\\r\\n\",\n \"g_data = g_data[...,[4,0,1,2,3]]\\r\\n\",\n \"g_data[...,1] = g_data[...,1] - g_data[...,3]/2\\r\\n\",\n \"g_data[...,2] = g_data[...,2] - g_data[...,4]/2\\r\\n\",\n \"\\r\\n\",\n \"g_data[0]\"\n ],\n \"outputs\": [\n {\n \"output_type\": \"execute_result\",\n \"data\": {\n \"text/plain\": [\n \"array([[0.37773186, 0.06993452, 0.19879478, 0.93849486, 0.37714368],\\n\",\n \" [0.63268024, 0.08424908, 0.09495369, 0.7752126 , 0.770834 ],\\n\",\n \" [0.31078017, 0.22171676, 0.65355074, 0.6689019 , 0.08199155],\\n\",\n \" [0.9312403 , 0.20145085, 0.0982542 , 0.66996616, 0.05629438],\\n\",\n \" [0.6019517 , 0.11892939, 0.76413304, 0.7681484 , 0.15887733],\\n\",\n \" [0.31345773, 0.22002116, 0.64820474, 0.67002124, 0.08382312],\\n\",\n \" [0.31812534, 0.22010046, 0.6488024 , 0.6710409 , 0.08644559],\\n\",\n \" [0.36010844, 0.07238191, 0.18851566, 0.9383161 , 0.40076703],\\n\",\n \" [0.93621445, 0.20135537, 0.09839028, 0.66991025, 0.05546013]],\\n\",\n \" dtype=float32)\"\n ]\n },\n \"metadata\": {\n \"tags\": []\n },\n \"execution_count\": 5\n }\n ],\n \"metadata\": {\n \"id\": \"u_RK-gO_YPYj\",\n \"colab\": {\n \"base_uri\": \"https://localhost:8080/\"\n },\n \"outputId\": \"a2453ea3-1302-44b8-dff5-67bfd9c979bc\"\n }\n },\n {\n \"cell_type\": \"markdown\",\n \"source\": [\n \"#Losses\\r\\n\",\n \"Overall three losses are calculated for the comparison:\\r\\n\",\n \"* Overlapping\\r\\n\",\n \"* Alignment\\r\\n\",\n \"* IoU\\r\\n\",\n \"\\r\\n\",\n \"The expressions for the calculations can be found in the Readme file shared.\"\n ],\n \"metadata\": {\n \"id\": \"YMkdEbASNVY6\"\n }\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 6,\n \"source\": [\n \"def overlapping_loss(result):\\r\\n\",\n \" losses=np.zeros(len(result))\\r\\n\",\n \" idx=0\\r\\n\",\n \" for i in result:\\r\\n\",\n \" over=0\\r\\n\",\n \" for j in range(len(i)):\\r\\n\",\n \" A=float(i[j][3]*i[j][4])\\r\\n\",\n \" if A==0:\\r\\n\",\n \" continue\\r\\n\",\n \" for k in range(len(i)):\\r\\n\",\n \" if j==k:\\r\\n\",\n \" continue\\r\\n\",\n \" x1=i[j][1]\\r\\n\",\n \" x2=i[j][1]+i[j][3]\\r\\n\",\n \" y1=i[j][2]\\r\\n\",\n \" y2=i[j][2]+i[j][4]\\r\\n\",\n \" x3=i[k][1]\\r\\n\",\n \" x4=i[k][1]+i[k][3]\\r\\n\",\n \" y3=i[k][2]\\r\\n\",\n \" y4=i[k][2]+i[k][4]\\r\\n\",\n \" x_over=max(min(x2,x4)-max(x1,x3),0)\\r\\n\",\n \" y_over=max(min(y2,y4)-max(y1,y3),0)\\r\\n\",\n \" over+=x_over*y_over/A\\r\\n\",\n \" losses[idx]=over\\r\\n\",\n \" idx+=1\\r\\n\",\n \" return np.mean(losses)*100\"\n ],\n \"outputs\": [],\n \"metadata\": {\n \"id\": \"MwXRP1cqv6qT\"\n }\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 7,\n \"source\": [\n \"def alignment_loss(result):\\r\\n\",\n \" xl =result[...,1] \\r\\n\",\n \" yl = result[...,2]\\r\\n\",\n \" \\r\\n\",\n \" xr = xl+result[...,3]\\r\\n\",\n \" yr = yl + result[...,4]\\r\\n\",\n \"\\r\\n\",\n \" xc = (xl + xr)/2\\r\\n\",\n \" yc = (yl + yr)/2\\r\\n\",\n \"\\r\\n\",\n \" ele = [xl , yl , xc, yc, xr, yr]\\r\\n\",\n \" ele1 = []\\r\\n\",\n \" epsilon = 0\\r\\n\",\n \" for element in ele:\\r\\n\",\n \" min_xl = np.ones(shape = element.shape)\\r\\n\",\n \" for i in range(len(element)):\\r\\n\",\n \" for j in range(len(element[i])):\\r\\n\",\n \" for k in range(len(element[i])): \\r\\n\",\n \" if j != k :\\r\\n\",\n \" min_xl[i][j] = min(min_xl[i][j],abs(element[i][j]-element[i][k])) \\r\\n\",\n \" min_xl = -np.log(1.0-min_xl + epsilon)\\r\\n\",\n \" ele1.append(min_xl)\\r\\n\",\n \" ele1 = np.min(np.array(ele1), axis = 0)\\r\\n\",\n \" ele1 = np.mean(np.sum(ele1 , axis = 1))\\r\\n\",\n \" return ele1*100\"\n ],\n \"outputs\": [],\n \"metadata\": {\n \"id\": \"vm1lD-2dHNtF\"\n }\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 8,\n \"source\": [\n \"def calculate_iou(result):\\r\\n\",\n \" losses=np.zeros(len(result))\\r\\n\",\n \" idx=0\\r\\n\",\n \" for i in result:\\r\\n\",\n \" iou=0\\r\\n\",\n \" for j in range(len(i)):\\r\\n\",\n \" for k in range(j+1,len(i)):\\r\\n\",\n \" x1=i[j][1]\\r\\n\",\n \" x2=i[j][1]+i[j][3]\\r\\n\",\n \" y1=i[j][2]\\r\\n\",\n \" y2=i[j][2]+i[j][4]\\r\\n\",\n \" x3=i[k][1]\\r\\n\",\n \" x4=i[k][1]+i[k][3]\\r\\n\",\n \" y3=i[k][2]\\r\\n\",\n \" y4=i[k][2]+i[k][4]\\r\\n\",\n \"\\r\\n\",\n \" box_1 = [[x1, y1], [x2, y1], [x2, y2], [x1, y2]]\\r\\n\",\n \" box_2 = [[x3, y3], [x4, y3], [x4, y4], [x3, y4]]\\r\\n\",\n \"\\r\\n\",\n \" poly_1 = Polygon(box_1)\\r\\n\",\n \" poly_2 = Polygon(box_2)\\r\\n\",\n \"\\r\\n\",\n \" if poly_1.union(poly_2).area!=0:\\r\\n\",\n \" iou += poly_1.intersection(poly_2).area / poly_1.union(poly_2).area\\r\\n\",\n \" losses[idx]=iou\\r\\n\",\n \" idx+=1\\r\\n\",\n \" return np.mean(losses)*100\"\n ],\n \"outputs\": [],\n \"metadata\": {\n \"id\": \"DHPkZGI1C9sc\"\n }\n },\n {\n \"cell_type\": \"markdown\",\n \"source\": [\n \"#Data Comparison\\n\",\n \"Calculation of metrics for original data of 1000 documents and for all models for 1000 documents (from unseen data).\"\n ],\n \"metadata\": {\n \"id\": \"dxJffkgwN6jC\"\n }\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 9,\n \"source\": [\n \"overlist = [overlapping_loss(publaydata[0:1000]), overlapping_loss(g_data[0:1000]), overlapping_loss(VAE_res[0:1000]),overlapping_loss(Transformer_res)]\\r\\n\",\n \"ioulist = [calculate_iou(publaydata[0:1000]), calculate_iou(g_data[0:1000]), calculate_iou(VAE_res[0:1000]),calculate_iou(Transformer_res)]\\r\\n\",\n \"alignlist = [alignment_loss(publaydata[0:1000]), alignment_loss(g_data[0:1000]), alignment_loss(VAE_res[0:1000]),alignment_loss(Transformer_res)]\"\n ],\n \"outputs\": [],\n \"metadata\": {\n \"id\": \"zkq1AtqOBDch\"\n }\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 10,\n \"source\": [\n \"#normalizing the lists w.r.t the original data\\r\\n\",\n \"overlist/=overlist[0]\\r\\n\",\n \"ioulist/=ioulist[0]\\r\\n\",\n \"alignlist/=alignlist[0]\\r\\n\",\n \"\\r\\n\",\n \"rows = [\\\"Original Data\\\", \\\"LayoutGAN\\\", \\\"LayoutVAE\\\",\\\"Layout Transformer\\\"]\\r\\n\",\n \"df=pd.DataFrame(index=rows)\\r\\n\",\n \"df[\\\"Overlap\\\"]=overlist\\r\\n\",\n \"df[\\\"IOU\\\"]=ioulist\\r\\n\",\n \"df[\\\"Alignment\\\"]=alignlist\"\n ],\n \"outputs\": [],\n \"metadata\": {\n \"id\": \"lumvx5zgC9yr\"\n }\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 11,\n \"source\": [\n \"print(\\\"Comparison for Publaynet Dataset\\\")\\r\\n\",\n \"display(df)\"\n ],\n \"outputs\": [\n {\n \"output_type\": \"stream\",\n \"name\": \"stdout\",\n \"text\": [\n \"Comparison for Publaynet Dataset\\n\"\n ]\n },\n {\n \"output_type\": \"display_data\",\n \"data\": {\n \"text/plain\": [\n \" Overlap IOU Alignment\\n\",\n \"Original Data 1.000000 1.000000 1.000000\\n\",\n \"LayoutGAN 1172.005234 2745.437529 1.164882\\n\",\n \"LayoutVAE 119.320127 185.864381 3.493406\\n\",\n \"Layout Transformer 1.090315 1.422297 0.739862\"\n ],\n \"text/html\": [\n \"
\\n\",\n \"\\n\",\n \"\\n\",\n \" \\n\",\n \" \\n\",\n \" \\n\",\n \" \\n\",\n \" \\n\",\n \" \\n\",\n \" \\n\",\n \" \\n\",\n \" \\n\",\n \" \\n\",\n \" \\n\",\n \" \\n\",\n \" \\n\",\n \" \\n\",\n \" \\n\",\n \" \\n\",\n \" \\n\",\n \" \\n\",\n \" \\n\",\n \" \\n\",\n \" \\n\",\n \" \\n\",\n \" \\n\",\n \" \\n\",\n \" \\n\",\n \" \\n\",\n \" \\n\",\n \" \\n\",\n \" \\n\",\n \" \\n\",\n \" \\n\",\n \" \\n\",\n \" \\n\",\n \" \\n\",\n \"
OverlapIOUAlignment
Original Data1.0000001.0000001.000000
LayoutGAN1172.0052342745.4375291.164882
LayoutVAE119.320127185.8643813.493406
Layout Transformer1.0903151.4222970.739862
\\n\",\n \"
\"\n ]\n },\n \"metadata\": {\n \"tags\": []\n }\n }\n ],\n \"metadata\": {\n \"colab\": {\n \"base_uri\": \"https://localhost:8080/\",\n \"height\": 190\n },\n \"id\": \"GRrfRwETFmF4\",\n \"outputId\": \"45f7104a-4558-43aa-d003-71b70489a259\"\n }\n }\n ]\n}" }, { "path": "README.md", "content": "# Layout Generation and Baseline Implementation\n\n## Contents\n* [Layout VAE](#layout-vae)\n * [Layout VAE Model](#layout-vae-model)\n * [Flow Diagram of Both Count and BBox VAE](#flow-diagram)\n * [Results Obtained](#results-obtained)\n* [Layout Transformer](#layout-transformer)\n * [Layout Transformer Model Architecture](#layout-transformer-model-architecture)\n * [Results](#results)\n* [LayoutGAN](#layoutgan)\n * [Architecture](#architecture)\n * [Results on MNIST](#results-on-mnist)\n * [Results on single column layouts](#results-on-single-column-layouts)\n* [Quantitative Comparison](#quantitative-comparison)\n\n\n## Layout VAE\nLayoutVAE is a variational autoencoder based model . It is a probabilistic and autoregressive model which generates the scene layout using latent variables in lower dimensions . It is capable of generating different layouts using the same data point.\n\n* **CountVAE:** This is the first part of the layoutVAE model; it takes the label set as input and predicts the counts of bounding boxes for corresponding labels. The input is provided as multilabel encoding.\n* **BBox VAE:** This the second part of the model was BBox VAE with LSTM based Embedding Generation. Similar to Countvae here also previous predictions along with the label set and label counts are used as conditioning info for current predictions.\n\n### Layout VAE Model \n![modelvae](https://user-images.githubusercontent.com/40228110/129761484-ba8b3494-67dc-437e-813e-705c9de19630.png)\n\n\n### Flow Diagram\n![Architecture](https://user-images.githubusercontent.com/40228110/129761516-a33098f9-15f1-4bcd-88de-04644beeae1c.png)\n\n\n### Results Obtained\n![VAE_result](/readme_images/VAE_result.png)\n\n## Layout Transformer\nLayout Transformer is a model proposed for generating structured layouts which can be used for documents, websites, apps, etc. It uses the decoder block of the Transformer Model, which is able to capture the relation of the document boxes with the previously predicted boxes (or inputs). Since it is an auto-regressive model, it can be used to generate entirely new layouts or to complete existing partial layouts.\nThe paper also emphasized on the fact that this model performs better than the existing models (at that time) and is better in the following aspects:\n* Able to generate layouts of arbitrary lengths\n* Gives better alignment due to the discretized grid\n* Is able to effectively capture the relationships between boxes in a single layout, which gives meaningful layouts\n\n### Layout Transformer Model Architecture\n![Trans_model](/readme_images/Trans_archi.png)\n\n### Results \n\n![Trans_result](/readme_images/Trans_res.png)\n\n## LayoutGAN\nLayoutGAN uses a GAN network , with the generator taking randomly sampled inputs (class probabilities and geometric parameters) as parameters, arranging them and thus producing refined geometric and class parameters.\n\n### Architecture \n\n\n### Results on MNIST\n![](LayoutGAN/demo/mnist_obtained.jpeg)\n\n### Results on single column layouts\n\n\n## Quantitative Comparison\nA total of three metrics were used to compare the models. \n* Overlapping Loss\n* Interection over Union (IoU)\n* Alignment Loss\n\nAfter Calculating the losses for each model, the following comparison table was obtained:\n\n| | Overlap | IOU | Alignment |\n|--------------------|:-----------:|:-----------:|:---------:|\n| Original Data | 1.000000 | 1.000000 | 1.000000 |\n| LayoutGAN | 1172.005234 | 2745.437529 | 1.164882 |\n| LayoutVAE | 119.320127 | 185.864381 | 3.493406 |\n| Layout Transformer | 1.090315 | 1.422297 | 0.739862 |\n" } ]