diff --git a/auto_encoder.py b/auto_encoder.py
index 54e58cb..11427cf 100644
--- a/auto_encoder.py
+++ b/auto_encoder.py
@@ -1,17 +1,18 @@
 import os
 import time
 import argparse
+from cv2 import polarToCart
 import wandb
-from cv2 import drawFrameAxes
+import random
 import numpy as np
 import tensorflow as tf
 from custom_utils.custom_utils import UV2Mesh, maskFacesBatch
 from tensorboardX import SummaryWriter
+import plotly.graph_objects as go
 
 from utils import str_to_bool
 from tf_utils import check_tf_version, allow_memory_growth, split_gpu_for_testing
 from load_models import load_generator, load_discriminator
-from dataset import DatasetSimple
 from losses import d_logistic, d_logistic_r1_reg, g_logistic_non_saturating, g_logistic_ns_pathreg, edge_loss, uv_mask_loss
 # from losses import collision_penalty as collision_penalty
 from losses import collision_loss as collision_penalty
@@ -19,8 +20,8 @@ from convnext.models.convnext_tf import ConvNeXt
 from stylegan2_ae.generator import Generator
 
 
-os.environ["CUDA_VISIBLE_DEVICES"] = "3"
-run = wandb.init(config=tf.compat.v1.flags.FLAGS, sync_tensorboard=True, project='stylegan', name='ae_1_2', notes='convnext + stylegan generator')
+os.environ["CUDA_VISIBLE_DEVICES"] = "2"
+run = wandb.init(config=tf.compat.v1.flags.FLAGS, sync_tensorboard=True, project='stylegan', name='ae_unposed_1', notes='convnext + stylegan generator')
 artifact = wandb.Artifact('code', type='code')
 artifact.add_file('auto_encoder.py')
 artifact.add_file('stylegan2_ae/layers/mapping_block.py')
@@ -29,6 +30,35 @@ run.log_artifact(artifact)
 uv2mesh = UV2Mesh(res=256)
 
 
+class DatasetSimple(tf.data.Dataset):
+
+    def _generator(folder):
+        # Opening the file
+        folder = folder.decode("utf-8")
+        uv_folder = os.path.join(folder, 'uv_posed') 
+        frames = os.listdir(uv_folder)
+        random.shuffle(frames)
+
+        for frame in frames:
+            id, frame_nr = frame.split('.')[0].split('_')
+            mesh_mask = np.load(os.path.join(folder, 'mesh_mask', id + '.npy')); mesh_mask = mesh_mask.astype(np.float32)
+            uv_mask = np.load(os.path.join(folder, 'uv_mask', id + '.npz')); uv_mask=uv_mask['uv_mask']; uv_mask = uv_mask.astype(np.float32)
+            mesh_unposed = np.load(os.path.join(folder, 'mesh_unposed', id + '_' + frame_nr + '.npz')); mesh_unposed = mesh_unposed['mesh_unposed']
+            uv_unposed = np.load(os.path.join(folder, 'uv_unposed', id + '_' + frame_nr + '.npz')); uv_unposed = uv_unposed['uv_unposed']; uv_unposed = np.transpose(uv_unposed, [2, 0, 1])
+
+            yield uv_unposed, mesh_unposed, uv_mask, mesh_mask
+
+    def __new__(cls, folder, res=256):
+        return tf.data.Dataset.from_generator(
+            cls._generator,
+            output_signature = (tf.TensorSpec(shape = (3, res * 2, res,), dtype = tf.float32), 
+                                tf.TensorSpec(shape = (14475, 3,), dtype = tf.float32), 
+                                tf.TensorSpec(shape = (res * 2, res,), dtype = tf.float32), 
+                                tf.TensorSpec(shape = (14475,), dtype = tf.float32), ),
+            args=(folder, )
+        )
+
+
 
 def filter_resolutions_featuremaps(resolutions, featuremaps, res):
     index = resolutions.index(res)
@@ -37,6 +67,43 @@ def filter_resolutions_featuremaps(resolutions, featuremaps, res):
     return filtered_resolutions, filtered_featuremaps
 
 
+def plot_3d(mesh_gt, mesh_pred, faces):
+    mesh_pred[:,1] += 2
+    fig = go.Figure(data=[
+        go.Mesh3d(
+            # 8 vertices of a cube
+            x=mesh_gt[:,0],
+            y=mesh_gt[:,1],
+            z=mesh_gt[:,2],
+            # i, j and k give the vertices of triangles
+            i = faces[:,0],
+            j = faces[:,1],
+            k = faces[:,2]
+        ),
+        go.Mesh3d(
+            # 8 vertices of a cube
+            x=mesh_pred[:,0],
+            y=mesh_pred[:,1],
+            z=mesh_pred[:,2],
+            # i, j and k give the vertices of triangles
+            i = faces[:,0],
+            j = faces[:,1],
+            k = faces[:,2],
+        )
+    ])
+
+    fig.update_layout(scene_aspectmode='manual',
+                    scene_aspectratio=dict(x=1, y=2, z=1))
+    fig.update_layout(
+        scene = dict(
+            xaxis = dict(nticks=4, range=[-1,1],),
+                        yaxis = dict(nticks=4, range=[-1,3],),
+                        zaxis = dict(nticks=4, range=[-1,1],),)
+                        )
+
+    return fig
+
+
 class AutoEncoder(tf.keras.Model):
     
     def __init__(self):
@@ -78,19 +145,18 @@ class AutoEncoder(tf.keras.Model):
 
 @tf.function
 def train_step(batch):
-    posed_uv, conditions,_, mesh_body_posed, mesh_mask = batch
-    uv_mask = conditions[3]
+    uv_unposed, mesh_unposed, uv_mask, mesh_mask = batch
 
     with tf.GradientTape() as tape:
 
-        pred = model([posed_uv, uv_mask])
+        pred = model([uv_unposed, uv_mask])
 
         mesh_pred = uv2mesh(pred)
-        mesh_orig = uv2mesh(posed_uv)
+        mesh_orig = mesh_unposed
 
         e_loss = 0.0#edge_loss(mesh_pred, mesh_orig, mesh_mask)
         c_loss = 0.0#collision_penalty(mesh_pred, mesh_body_posed, mesh_mask)
-        l1_loss = tf.reduce_mean(tf.reduce_sum((tf.math.abs(posed_uv - pred)), axis=1) * uv_mask)
+        l1_loss = tf.reduce_mean(tf.reduce_sum((tf.math.abs(uv_unposed - pred)), axis=1) * uv_mask)
         loss = l1_loss + e_loss + c_loss
 
     metric = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(mesh_mask[:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(mesh_mask)
@@ -102,16 +168,15 @@ def train_step(batch):
 
 @tf.function
 def valid_step(batch):
-    posed_uv, conditions,_, mesh_body_posed, mesh_mask = batch
-    uv_mask = conditions[3]
+    uv_unposed, mesh_unposed, uv_mask, mesh_mask = batch
 
-    pred = model([posed_uv, uv_mask])
+    pred = model([uv_unposed, uv_mask])
     mesh_pred = uv2mesh(pred)
-    mesh_orig = uv2mesh(posed_uv)
+    mesh_orig = mesh_unposed
 
     e_loss = 0.0#edge_loss(mesh_pred, mesh_orig, mesh_mask)
     c_loss = 0.0#collision_penalty(mesh_pred, mesh_body_posed, mesh_mask)
-    l1_loss = tf.reduce_mean(tf.reduce_sum((tf.math.abs(posed_uv - pred)), axis=1) * uv_mask)
+    l1_loss = tf.reduce_mean(tf.reduce_sum((tf.math.abs(uv_unposed - pred)), axis=1) * uv_mask)
     loss = l1_loss + e_loss + c_loss
 
     metric = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(mesh_mask[:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(mesh_mask)
@@ -122,7 +187,7 @@ if __name__ == '__main__':
     allow_memory_growth()
     
     
-    name = 'test-9'
+    name = 'ae_unposed_1'
     dataset_train = DatasetSimple('/home/workspace/data/dataset/stylegan_new/train').cache().batch(8).repeat()
     dataset_valid = DatasetSimple('/home/workspace/data/dataset/stylegan_new/valid').cache().batch(8).repeat()
 
@@ -131,8 +196,7 @@ if __name__ == '__main__':
     optimizer = tf.optimizers.Adam()
     folder = 'models/ae_logs/'
     train_summary_writer = tf.summary.create_file_writer(folder + name + '/train')
-    valid_summary_writer = tf.summary.create_file_writer(folder + name + '/valid')
-    mesh_writer = SummaryWriter(folder + name + '/mesh')
+    # valid_summary_writer = tf.summary.create_file_writer(folder + name + '/valid')
 
     step = 0
     for batch in dataset_train:
@@ -142,18 +206,19 @@ if __name__ == '__main__':
             with train_summary_writer.as_default():
                 tf.summary.scalar(key, train_loss[key], step=step)
 
-        data_valid = iter(dataset_valid).get_next()
-        valid_loss = valid_step(data_valid)
+        data_valid = batch
+        # data_valid = iter(dataset_valid).get_next()
+        # valid_loss = valid_step(data_valid)
 
-        for key in valid_loss.keys():
-            with valid_summary_writer.as_default():
-                tf.summary.scalar(key, valid_loss[key], step=step)
+        # for key in valid_loss.keys():
+        #     with valid_summary_writer.as_default():
+        #         tf.summary.scalar(key, valid_loss[key], step=step)
 
         if step % 250 == 0:
-            pred = model([data_valid[0], data_valid[1][3]])
+            pred = model([data_valid[0], data_valid[2]])
             
             mesh_pred = uv2mesh(pred)
-            mesh_orig = uv2mesh(data_valid[0])
+            mesh_orig = data_valid[1]
             
             orig_img = tf.transpose(data_valid[0], [0,2,3,1])
             pred_img = tf.transpose(pred, [0,2,3,1])
@@ -162,12 +227,13 @@ if __name__ == '__main__':
                 tf.summary.image('images', img, step=step)
 
             F = np.load("custom_utils/smpl_faces.npy")
-            garment, faces = maskFacesBatch(mesh_pred.numpy(), F, data_valid[4].numpy().astype(bool))
-            gt, _ = maskFacesBatch(mesh_orig.numpy(), F, data_valid[4].numpy().astype(bool))
+            garment, faces = maskFacesBatch(mesh_pred.numpy(), F, data_valid[3].numpy().astype(bool))
+            gt = [mesh_orig[i].numpy()[data_valid[3][i].numpy().astype(bool), :] for i in range(pred.shape[0])]
             
-            for i in range(len(garment)):
-                mesh_writer.add_mesh('garment_pred_' + str(i), vertices=garment[i][np.newaxis,:,:], faces=faces[i][np.newaxis,:,:], global_step=step)
-                mesh_writer.add_mesh('garment_gt_' + str(i), vertices=gt[i][np.newaxis,:,:], faces=faces[i][np.newaxis,:,:], global_step=step)
+            for i in range(pred.shape[0]):
+                fig = plot_3d(gt[i], garment[i], faces[i])
+                wandb.log({"garment_" + str(i): fig})
+                
 
         print(step)
 
diff --git a/custom_utils/__init__.py b/custom_utils/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/custom_utils/graph_utils.py b/custom_utils/graph_utils.py
new file mode 100644
index 0000000..5e09d28
--- /dev/null
+++ b/custom_utils/graph_utils.py
@@ -0,0 +1,207 @@
+import scipy as sp
+import numpy as np
+import scipy.io as sio
+import tensorflow as tf
+from tensorflow.keras.activations import tanh, relu
+from tensorflow.keras.layers import Dense, BatchNormalization, Add, Flatten
+from custom_utils.custom_utils import UV2Mesh
+
+from tensorflow_graphics.geometry.convolution.graph_pooling import pool, unpool
+from tensorflow_graphics.geometry.convolution.graph_convolution import edge_convolution_template
+
+
+def sparse_to_tensor(L):
+    idx = np.mat([L.row, L.col]).transpose().astype(np.float32)
+    data = L.data.astype(dtype=np.float32)
+    shape = L.shape
+    return tf.SparseTensor(idx, data, shape)
+
+def merge_verts(Vs, batchDim):
+    V = tf.concat(tf.unstack(Vs, num=batchDim), axis=0)
+    return V
+
+
+def repeat_laplacian(L, times):
+    idx = np.repeat(np.mat([L.row, L.col]), repeats=times, axis=1)
+    idx_offset = np.linspace(0, L.shape[0]  * (times-1), times).astype(np.int32)
+    idx_offset = np.repeat(idx_offset, L.row.shape[0], axis=0)
+    idx += idx_offset[np.newaxis,:]
+    data = np.repeat(L.data, repeats=times, axis=0)
+    shape = (L.shape[0] * times, L.shape[1] * times)
+    return sp.sparse.coo_matrix((data, (np.squeeze(np.asarray(idx[0,:])), np.squeeze(np.asarray(idx[1,:])))), shape)
+
+
+
+class GraphConvolution(tf.keras.layers.Layer):
+    def __init__(self, shape, activation=None, name='gcn', L=None, batchSize=None, normalization=None):
+        super(GraphConvolution, self).__init__()
+        self.shape = shape
+        self.w = self.add_weight(name=name + '/weight', shape=(2, *shape), initializer=tf.initializers.glorot_normal, trainable=True)
+        self.b = self.add_weight(name=name + '/bias', shape=shape[-1], initializer=tf.initializers.zeros, trainable=True)
+        self.activation = activation or (lambda x: x)
+        self.normalization = normalization
+        self.batchSize = batchSize
+        self.indices = tf.cast(tf.linspace(0, self.batchSize * L.shape[1], self.batchSize+1), dtype=tf.int32)
+        self.L = repeat_laplacian(L, self.batchSize)
+
+    def gather(self):
+        return [self.w, self.b]
+
+    def _single(self, X):
+        X0 = tf.einsum('ab,bc->ac', X, self.w[0]) # Node
+        X1 = tf.einsum('ab,bc->ac', X, self.w[1]) # Neigh
+        # Neigh conv
+        X1 = edge_convolution_template(
+            X1, sparse_to_tensor(self.L), sizes=None,
+            edge_function=lambda x, y: y, 
+            reduction='weighted', 
+            edge_function_kwargs={}
+        ) 
+        X = X0 + X1 + self.b
+        return X
+
+    def call(self, X, training):
+        X = merge_verts(X, self.batchSize)
+        X = self._single(X)
+        X = tf.stack([X[self.indices[i-1]:self.indices[i]] for i in range(1, len(self.indices))])
+        X = self.activation(X)
+        if self.normalization:
+            X = self.normalization(X, training=training)
+        return X
+
+
+class GraphUnPool(tf.keras.layers.Layer):
+    def __init__(self, pool_map):
+        super(GraphUnPool, self).__init__()
+        self.pool_map = pool_map
+
+    def call(self, X):
+        # print(X.dtype)
+        # print(self.pool_map.dtype)
+        return unpool(X, self.pool_map, None)
+        # return tf.gather(X, self.pool_map, axis=1)
+
+
+
+class GraphPool(tf.keras.layers.Layer):
+    def __init__(self, pool_map, type, batch_size):
+        super(GraphPool, self).__init__()
+        self.type = type
+        self.batch_size = batch_size
+        self.indices = tf.cast(tf.linspace(0, batch_size * pool_map.shape[0], batch_size+1), dtype=tf.int32)
+        # self.pool_map = repeat_laplacian(pool_map, self.batch_size)
+        self.pool_map = np.squeeze(pool_map)
+
+    def _single(self, x):
+        return pool(x, sparse_to_tensor(self.pool_map), algorithm=self.type, sizes=None)
+
+    def call(self, x, training=None):
+        # x = merge_verts(x, self.batch_size)
+        # x = self._single(x)
+        # x = tf.stack([x[self.indices[i-1]:self.indices[i]] for i in range(1, len(self.indices))])
+
+        N, M = len(self.pool_map), len(set(self.pool_map)) # N = HQ verts, M = LQ verts
+        # Assign correspondences as V0_lq = [V0_hq, V1_hq, ..., VN_hq]
+        d = [[] for _ in range(M)] # list of lists
+        for i in range(N): d[self.pool_map[i]] += [i]
+        # max sublist len
+        m = max([len(d[i]) for i in range(M)])
+        # expand all sublist to max len
+        for i in range(M):
+            while len(d[i]) < m: 
+                # repeating elements will NOT change MIN/MAX output
+                d[i] += [d[i][0]] 
+ 
+        x = tf.reduce_max(tf.gather(x, d, axis=1), axis=2)
+        return x
+
+
+
+
+
+class GNNDiscriminator(tf.keras.Model):
+    def __init__(self, batchSize, isSkirt=False, n_feature=256, pool_type="max", global_average=True):
+        super(GNNDiscriminator, self).__init__()
+
+        self.uv2mesh = UV2Mesh(res=256)        
+
+        laplacians = [sio.loadmat("custom_utils/laplacian/ssmpl" + ("_skirt" if isSkirt else "") + str(i) + ".mat")['L'].tocoo() for i in range(5)] 
+        poolings = [sio.loadmat("custom_utils/pooling/pool" + str(i) + str(i+1) + ".mat")['idx'] for i in range(5)]
+   
+        self.conv1 = GraphConvolution((3, int(n_feature/2**2)), activation=relu, L=laplacians[0], batchSize=batchSize, normalization=BatchNormalization(), name='gcn_1')
+        self.pool1 = GraphPool(pool_map=poolings[0], type=pool_type, batch_size=batchSize)
+
+        self.conv2 = GraphConvolution((int(n_feature/2**2), int(n_feature/2**2)), activation=relu, L=laplacians[1], batchSize=batchSize, normalization=BatchNormalization(), name='gcn_2')
+        self.pool2 = GraphPool(pool_map=poolings[1], type=pool_type, batch_size=batchSize)
+
+        self.custom_layers = []
+        for i in range(2,5):
+            self.custom_layers.append(GraphConvolution((int(n_feature/2**(4- (i if i == 2 else i-1) )), int(n_feature/2**(4-i))), activation=relu, L=laplacians[i], batchSize=batchSize, normalization=BatchNormalization()))
+            for j in range(2**(i-2)):
+                conv1 = GraphConvolution((int(n_feature/2**(4-i)), int(n_feature/2**(4-i))), activation=relu, L=laplacians[i], batchSize=batchSize, normalization=BatchNormalization())
+                conv2 = GraphConvolution((int(n_feature/2**(4-i)), int(n_feature/2**(4-i))), activation=relu, L=laplacians[i], batchSize=batchSize, normalization=BatchNormalization())
+                self.custom_layers.append([conv1, conv2])
+            self.custom_layers.append(GraphPool(pool_map=poolings[i], type=pool_type, batch_size=batchSize))
+            
+            self.custom_layers.append(Dense(256, relu))
+            self.custom_layers.append(BatchNormalization())
+            self.custom_layers.append(Dense(int(n_feature/2**(4-i)), relu))
+            self.custom_layers.append(BatchNormalization())
+
+        self.flatten = Flatten()
+        self.last = Dense(1)
+
+    def call(self, inputs, training=None):
+        inputs = self.uv2mesh(inputs)
+
+        x = self.conv1(inputs, training)
+        x = self.pool1(x)
+
+        x = self.conv2(x, training)
+        x = self.pool2(x)
+
+        for layer in self.custom_layers:
+            if isinstance(layer, list):
+                x1 = layer[0](x, training)
+                x2 = layer[1](x1, training)
+                x = Add()([x, x2])
+            else:
+                x = layer(x, training=training)
+
+        x = self.flatten(x)
+        x = self.last(x)
+
+        return x
+
+
+
+
+
+
+class DatasetTest(tf.data.Dataset):
+
+    def _generator():
+        for i in range(10):
+            yield np.array([i], dtype=np.int32)
+
+    def __new__(cls):
+        return tf.data.Dataset.from_generator(
+            cls._generator,
+            output_signature = tf.TensorSpec(shape = (1), dtype = tf.int32) ,
+            args=()
+        )
+
+
+
+if __name__ == '__main__':
+    import os
+    os.environ["CUDA_VISIBLE_DEVICES"] = "1"
+
+    batch_size = 4
+    input = tf.ones((batch_size, 14475, 3))
+
+    ge = GNNDiscriminator(batchSize=batch_size, n_feature=32)
+
+    x = ge(input)
+
+    print()
\ No newline at end of file
diff --git a/custom_utils/laplacian/ssmpl0.mat b/custom_utils/laplacian/ssmpl0.mat
new file mode 100644
index 0000000..e3cea72
Binary files /dev/null and b/custom_utils/laplacian/ssmpl0.mat differ
diff --git a/custom_utils/laplacian/ssmpl1.mat b/custom_utils/laplacian/ssmpl1.mat
new file mode 100644
index 0000000..5fe7629
Binary files /dev/null and b/custom_utils/laplacian/ssmpl1.mat differ
diff --git a/custom_utils/laplacian/ssmpl2.mat b/custom_utils/laplacian/ssmpl2.mat
new file mode 100644
index 0000000..71581f4
Binary files /dev/null and b/custom_utils/laplacian/ssmpl2.mat differ
diff --git a/custom_utils/laplacian/ssmpl3.mat b/custom_utils/laplacian/ssmpl3.mat
new file mode 100644
index 0000000..af9fcbf
Binary files /dev/null and b/custom_utils/laplacian/ssmpl3.mat differ
diff --git a/custom_utils/laplacian/ssmpl4.mat b/custom_utils/laplacian/ssmpl4.mat
new file mode 100644
index 0000000..98a8971
Binary files /dev/null and b/custom_utils/laplacian/ssmpl4.mat differ
diff --git a/custom_utils/laplacian/ssmpl5.mat b/custom_utils/laplacian/ssmpl5.mat
new file mode 100644
index 0000000..85d82b6
Binary files /dev/null and b/custom_utils/laplacian/ssmpl5.mat differ
diff --git a/custom_utils/laplacian/ssmpl_skirt0.mat b/custom_utils/laplacian/ssmpl_skirt0.mat
new file mode 100644
index 0000000..3755659
Binary files /dev/null and b/custom_utils/laplacian/ssmpl_skirt0.mat differ
diff --git a/custom_utils/laplacian/ssmpl_skirt1.mat b/custom_utils/laplacian/ssmpl_skirt1.mat
new file mode 100644
index 0000000..501181e
Binary files /dev/null and b/custom_utils/laplacian/ssmpl_skirt1.mat differ
diff --git a/custom_utils/laplacian/ssmpl_skirt2.mat b/custom_utils/laplacian/ssmpl_skirt2.mat
new file mode 100644
index 0000000..8e7b579
Binary files /dev/null and b/custom_utils/laplacian/ssmpl_skirt2.mat differ
diff --git a/custom_utils/laplacian/ssmpl_skirt3.mat b/custom_utils/laplacian/ssmpl_skirt3.mat
new file mode 100644
index 0000000..6b098ff
Binary files /dev/null and b/custom_utils/laplacian/ssmpl_skirt3.mat differ
diff --git a/custom_utils/laplacian/ssmpl_skirt4.mat b/custom_utils/laplacian/ssmpl_skirt4.mat
new file mode 100644
index 0000000..16b677b
Binary files /dev/null and b/custom_utils/laplacian/ssmpl_skirt4.mat differ
diff --git a/custom_utils/laplacian/ssmpl_skirt5.mat b/custom_utils/laplacian/ssmpl_skirt5.mat
new file mode 100644
index 0000000..028b425
Binary files /dev/null and b/custom_utils/laplacian/ssmpl_skirt5.mat differ
diff --git a/custom_utils/pooling/pool01.mat b/custom_utils/pooling/pool01.mat
new file mode 100644
index 0000000..a48d290
Binary files /dev/null and b/custom_utils/pooling/pool01.mat differ
diff --git a/custom_utils/pooling/pool12.mat b/custom_utils/pooling/pool12.mat
new file mode 100644
index 0000000..f70697c
Binary files /dev/null and b/custom_utils/pooling/pool12.mat differ
diff --git a/custom_utils/pooling/pool23.mat b/custom_utils/pooling/pool23.mat
new file mode 100644
index 0000000..60f89db
Binary files /dev/null and b/custom_utils/pooling/pool23.mat differ
diff --git a/custom_utils/pooling/pool34.mat b/custom_utils/pooling/pool34.mat
new file mode 100644
index 0000000..f694339
Binary files /dev/null and b/custom_utils/pooling/pool34.mat differ
diff --git a/custom_utils/pooling/pool45.mat b/custom_utils/pooling/pool45.mat
new file mode 100644
index 0000000..f241ab3
Binary files /dev/null and b/custom_utils/pooling/pool45.mat differ
diff --git a/custom_utils/smpl.py b/custom_utils/smpl.py
index a6f38eb..ea87064 100644
--- a/custom_utils/smpl.py
+++ b/custom_utils/smpl.py
@@ -121,6 +121,7 @@ class SMPL:
 
     def run(self, R, betas, remove=True):
         batch_size = betas.shape[0]
+        # R = self.run_rodrigues(pose)
 
         #J_regressor = tf.gather(self.J_regressor, gender, axis=0)
         #shapedirs = tf.gather(self.shapedirs, gender, axis=0)
diff --git a/custom_utils/uv_zero_body.npy b/custom_utils/uv_zero_body.npy
new file mode 100644
index 0000000..a231698
Binary files /dev/null and b/custom_utils/uv_zero_body.npy differ
diff --git a/load_models.py b/load_models.py
index 2ab2359..6fdae31 100644
--- a/load_models.py
+++ b/load_models.py
@@ -38,23 +38,20 @@ def load_generator(g_params=None, is_g_clone=False, ckpt_dir=None, custom_cuda=T
 
 
 def load_discriminator(d_params=None, ckpt_dir=None, custom_cuda=True):
-    if custom_cuda:
-        from stylegan2.discriminator import Discriminator
-    else:
-        from stylegan2_ref.discriminator import Discriminator
+    # if custom_cuda:
+    #     from stylegan2.discriminator import Discriminator
+    # else:
+    #     from stylegan2_ref.discriminator import Discriminator
+
+    from custom_utils.graph_utils import GNNDiscriminator as Discriminator
 
-    if d_params is None:
-        d_params = {
-            'labels_dim': 0,
-            'resolutions': [4, 8, 16, 32, 64, 128, 256, 512, 1024],
-            'featuremaps': [512, 512, 512, 512, 512, 256, 128, 64, 32],
-        }
 
     res = d_params['resolutions'][-1]
-    test_images = tf.ones((1, 3, res, int(res/2)), dtype=tf.float32)
+    batch_size = d_params['batch_size']
+    test_images = tf.ones((batch_size, 3, res, int(res/2)), dtype=tf.float32)
 
     # build discriminator model
-    discriminator = Discriminator(d_params)
+    discriminator = Discriminator(batch_size, 32)
     _ = discriminator(test_images)
 
     if ckpt_dir is not None:
diff --git a/stylegan2/generator.py b/stylegan2/generator.py
index b6e9515..2bcc193 100644
--- a/stylegan2/generator.py
+++ b/stylegan2/generator.py
@@ -17,6 +17,8 @@ class Generator(tf.keras.Model):
         self.featuremaps = g_params['featuremaps']
         self.w_ema_decay = 0.995
         self.style_mixing_prob = 0.9
+        self.zero_body = tf.ones((4,3,512,256)) * tf.transpose(g_params['zero_body'], [2,0,1])[tf.newaxis,:,:,:]
+
 
         self.n_broadcast = len(self.resolutions) * 2
         self.mixing_layer_indices = tf.range(self.n_broadcast, dtype=tf.int32)[tf.newaxis, :, tf.newaxis]
@@ -102,7 +104,8 @@ class Generator(tf.keras.Model):
             w_broadcasted = self.truncation_trick(w_broadcasted, truncation_psi, truncation_cutoff)
 
         image_out = self.synthesis(w_broadcasted)
-        # image_out = image_out + conditions[4]
+        # image_out = image_out + conditions[2]
+        image_out = image_out + self.zero_body[:image_out.shape[0]]
         image_out = image_out * conditions[3][:,tf.newaxis,:,:]
 
         if ret_w_broadcasted:
diff --git a/train.py b/train.py
index 1e3709a..3d12fb7 100644
--- a/train.py
+++ b/train.py
@@ -18,7 +18,7 @@ from losses import collision_loss as collision_penalty
 from custom_utils.smpl import SMPL
 
 os.environ["CUDA_VISIBLE_DEVICES"] = "0"
-run = wandb.init(config=tf.compat.v1.flags.FLAGS, sync_tensorboard=True, project='stylegan', name='run_28', notes='learning offsets, all losses from the start, smaller generator (from 128), slower learning rate for generator, add a fading out l1 loss with gt, add tanh activation for last layer, higher weights for the extra losses, same learning rate again. added extra metrics, lower discr. features too. use reduce_sum for losses divided by sum of mask. learn unposed version then do skinning. lowered collsion coeff')
+run = wandb.init(config=tf.compat.v1.flags.FLAGS, sync_tensorboard=True, project='stylegan', name='run_33', notes='learning offsets, all losses from the start, smaller generator (from 128), slower learning rate for generator, add a fading out l1 loss with gt, add tanh activation for last layer, higher weights for the extra losses, same learning rate again. added extra metrics, lower discr. features too. use reduce_sum for losses divided by sum of mask. learn unposed version then do skinning. lowered collsion coeff. learns offsets over zero body. fixed unposed uv map conversion. gnn discriminator.')
 artifact = wandb.Artifact('code', type='code')
 artifact.add_file('train.py')
 artifact.add_file('losses.py')
@@ -579,8 +579,10 @@ def main():
         'n_mapping': 8,
         'resolutions': train_resolutions,
         'featuremaps': train_g_featuremaps,
+        'zero_body': np.load('custom_utils/uv_zero_body.npy')
     }
     d_params = {
+        'batch_size': args['batch_size_per_replica'],
         'labels_dim': -1,
         'resolutions': train_resolutions,
         'featuremaps': train_d_featuremaps,
@@ -592,7 +594,7 @@ def main():
 
     # prepare dataset
     # dataset = get_ffhq_dataset(args['tfrecord_dir'], args['train_res'], batch_size=global_batch_size, epochs=None)
-    dataset = DatasetSimple('/home/workspace/data/dataset/stylegan_new/train').cache().batch(global_batch_size).repeat()
+    dataset = DatasetSimple('/home/workspace/data/dataset/stylegan_new/train').cache().batch(global_batch_size, drop_remainder=True).repeat()
 
     options = tf.data.Options()
     options.experimental_distribute.auto_shard_policy = tf.data.experimental.AutoShardPolicy.OFF
@@ -625,7 +627,7 @@ def main():
             'train_res': args['train_res'],
         }
 
-        trainer = Trainer(training_parameters, name=f'stylegan2-test-28-{args["train_res"]}x{args["train_res"]}')
+        trainer = Trainer(training_parameters, name=f'stylegan2-test-33-{args["train_res"]}x{args["train_res"]}')
         trainer.train(dist_dataset, strategy)
     return