run-35

Dockerfile

@@ -1,8 +0,0 @@
-FROM tensorflow/tensorflow:latest-gpu
-
-RUN apt-get update
-RUN apt-get upgrade -y
-RUN apt-get install git ffmpeg libsm6 libxext6 libopenexr-dev zip unzip -y
-RUN pip install -r requierements.txt
-
-# RUN wandb login

custom_utils/custom_utils.py

@@ -136,3 +136,15 @@ def plot_3d(mesh_gt, mesh_pred_unposed, faces, collision, collision_faces, colli
                         )
 
     return fig
+
+
+def batch_to_dict(batch, signature):
+    result = dict()
+    for i, key in enumerate(signature.keys()):
+            result[key] = batch[i]
+    return result
+
+conditions = ['pose_shape', 'uv_static', 'uv_mask', 'uv_shaped_body']
+def nest_conditions(batch):
+    batch['conditions'] = {d: batch.pop(d) for d in conditions}
+    return batch

dataset.py

@@ -6,6 +6,10 @@ import random
 import numpy as np
 import tensorflow as tf
 
+from custom_utils.custom_utils import nest_conditions
+
+
+
 class DatasetSimple(tf.data.Dataset):
 
     def _generator(folder):
@@ -47,10 +51,182 @@ class DatasetSimple(tf.data.Dataset):
             args=(folder, )
         )
 
-os.environ["CUDA_VISIBLE_DEVICES"] = "1"
+
+
+class StaticDataset():
+    def __init__(self, folder, signature) -> None:
+        self.store = dict()
+        self.signature = signature
+        self.folder = folder
+        self.possible_keys = ['uv_static', 'mesh_static', 'mesh_mask', 'uv_mask', 'uv_shaped_body']
+        self.keys = [value for value in signature.keys() if value in self.possible_keys]
+
+    def read(self, id):
+        data = dict()
+
+        if 'uv_static' in self.signature.keys():
+            uv_static = np.load(os.path.join(self.folder, 'uv_static', id +  '.npz')); uv_static = uv_static['uv_static']; uv_static = np.transpose(uv_static, [2, 0, 1]); data['uv_static']=uv_static
+        if 'mesh_static' in self.signature.keys():
+            mesh_static = np.load(os.path.join(self.folder, 'mesh_static', id +  '.npz')); mesh_static = mesh_static['mesh_static']; data['mesh_static'] = mesh_static
+        if 'mesh_mask' in self.signature.keys():
+            mesh_mask = np.load(os.path.join(self.folder, 'mesh_mask', id + '.npy')); mesh_mask = mesh_mask.astype(np.float32); data['mesh_mask'] = mesh_mask
+        if 'uv_mask' in self.signature.keys():
+            uv_mask = np.load(os.path.join(self.folder, 'uv_mask', id + '.npz')); uv_mask=uv_mask['uv_mask']; uv_mask = uv_mask.astype(np.float32); data['uv_mask'] = uv_mask
+        if 'uv_shaped_body' in self.signature.keys():
+            uv_shaped_body = np.load(os.path.join(self.folder, 'uv_shaped_body', id + '.npz')); uv_shaped_body = uv_shaped_body['uv_shaped_body']; uv_shaped_body = np.transpose(uv_shaped_body, [2, 0, 1]); data['uv_shaped_body'] = uv_shaped_body
+
+        return data
+
+    def __getitem__(self, id) -> dict:
+        # id = tf.strings.unicode_decode(id, 'UTF-8')
+        # id = tf.strings.as_string(id)
+        # id = id.numpy()
+        if id not in self.store:
+            self.store[id] = self.read(id)
+
+        return self.store[id]
+
+    def expand_batch(self, batch):
+        for key in self.keys:
+                batch[key] = tf.zeros(((0,) + self.signature[key]))
+        for i in range(batch['id'].shape[0]):
+            item = self.__getitem__(batch['id'][i].numpy().decode('UTF-8'))
+            for key in self.keys:
+                batch[key] = tf.concat((batch[key], tf.expand_dims(item[key], axis=0)), axis=0)
+        return batch
+
+
+
+class DatasetIdx(tf.data.Dataset):
+
+    def __new__(cls, folder):
+        uv_static_folder = os.path.join(folder, 'pose_shape') 
+        frames = os.listdir(uv_static_folder)
+        frames = [frame.split('.')[0] for frame in frames] # .split('_')[0]
+        random.shuffle(frames)
+
+        return tf.data.Dataset.from_tensor_slices(frames)
+
+
+# class DatasetFactory(tf.data.Dataset):
+class DatasetFactory():
+
+    def __init__(self):
+        super(DatasetFactory, self).__init__()
+
+    def _build_tf_signature(self, signature):
+        tf_signature = []
+        for key in signature.keys():
+            if key == 'id':
+                tf_signature.append(tf.TensorSpec(shape = (), dtype = tf.string))
+            else:
+                tf_signature.append(tf.TensorSpec(shape = signature[key], dtype = tf.float32))
+        return tf_signature
+
+    
+    def get_dataset(self, folder, batch_size, signature):
+        # static_data = StaticDataset(folder, signature)
+        dataset_idx = DatasetIdx(folder)
+
+        def read_item(frame):
+            frame_split = tf.strings.split(frame, sep='_')
+            id, frame_nr = frame_split[0].numpy().decode("utf-8"), frame_split[1].numpy().decode("utf-8")
+            # id, frame_nr = frame_split[0].numpy(), frame_split[1].numpy()
+            # static_info = static_data[id]
+            item = dict()
+
+            if 'id' in signature.keys():
+                item['id'] = id
+            if 'pose_shape' in signature.keys():
+                pose_shape = np.load(os.path.join(folder, 'pose_shape',  id + '_' + frame_nr + '.npy'))
+                item['pose_shape'] = pose_shape
+            if 'uv_posed' in signature.keys():
+                uv_posed = np.load(os.path.join(folder, 'uv_posed', id + '_' + frame_nr + '.npz')); uv_posed = uv_posed['uv_gt']; uv_posed = np.transpose(uv_posed, [2, 0, 1])
+                item['uv_posed'] = uv_posed
+            if 'uv_body_posed' in signature.keys():
+                uv_body_posed = np.load(os.path.join(folder, 'uv_body_posed', id + '_' + frame_nr + '.npz')); uv_body_posed = uv_body_posed['uv_body']; uv_body_posed = np.transpose(uv_body_posed, [2, 0, 1])
+                item['uv_body_posed'] = uv_body_posed
+            if 'mesh_body_posed' in signature.keys():
+                mesh_body_posed = np.load(os.path.join(folder, 'mesh_body_posed', id + '_' + frame_nr + '.npy'))
+                item['mesh_body_posed'] = mesh_body_posed
+            if 'mesh_posed' in signature.keys():
+                mesh_posed = np.load(os.path.join(folder, 'mesh_posed', id + '_' + frame_nr + '.npz')); mesh_posed = mesh_posed['mesh_posed']
+                item['mesh_posed'] = mesh_posed
+            if 'mesh_unposed' in signature.keys():
+                mesh_unposed = np.load(os.path.join(folder, 'mesh_unposed', id + '_' + frame_nr + '.npz')); mesh_unposed = mesh_unposed['mesh_unposed']
+                item['mesh_unposed'] = mesh_unposed
+            if 'uv_unposed' in signature.keys():
+                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])
+                item['uv_unposed'] = uv_unposed
+
+
+            # item = { **static_info, **item }
+            as_list = []
+            for key in signature.keys():
+                if key in item:
+                    as_list.append(item[key])
+
+            return as_list
+
+        dataset = dataset_idx.map(lambda i: tf.py_function(func=read_item, 
+                                               inp=[i], 
+                                               Tout=self._build_tf_signature(signature)
+                                               ), num_parallel_calls=tf.data.AUTOTUNE).cache().batch(batch_size=batch_size, drop_remainder=True).repeat()
+        # dataset = dataset_idx.map(read_item, num_parallel_calls=tf.data.AUTOTUNE).cache().batch(batch_size=batch_size, drop_remainder=True).repeat()
+
+        return dataset
+
+    # def __new__(cls, folder):
+    #     return tf.data.Dataset.from_generator(
+    #         cls._generator,
+    #         # output_signature = (tf.TensorSpec(shape = (), dtype = tf.string), ),
+    #         output_types=(tf.string),
+    #         args=(folder, )
+    #     )
+
+
+
+os.environ["CUDA_VISIBLE_DEVICES"] = "3"
 
 
 if __name__ == '__main__':
+    from custom_utils.custom_utils import batch_to_dict
+
+    factory = DatasetFactory()
+    signature_dynamic = { 'id': (),
+                        #   'pose_shape': (72 + 11,),
+                        #   'uv_posed': (3, 512, 256,),
+                        #   'uv_body_posed': (3, 512, 256,),
+                          'mesh_body_posed': (14475, 3,),
+                          'mesh_posed': (14475, 3,),
+                        #   'mesh_unposed': (14475, 3,),
+                        #   'uv_unposed': (3, 512, 256,),
+                        }
+    signature_static = { #'uv_static': (3, 512, 256,),
+                        #  'mesh_static': (14475, 3,),
+                         'mesh_mask': (14475,),
+                        #  'uv_mask': (512, 256,),
+                        #  'uv_shaped_body': (3, 512, 256,),
+                          }
+
+    dataset = factory.get_dataset(folder='/home/workspace/data/dataset/stylegan_new/train',
+                                  batch_size=1024,
+                                  signature=signature_dynamic)
+    static_dataset = StaticDataset(folder='/home/workspace/data/dataset/stylegan_new/train',
+                                   signature=signature_static)
+
+    @tf.function
+    def euclid_error(mask, gt, pred):
+        return tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(mask[:,:,tf.newaxis] * (gt - pred) ** 2, axis=-1))) / tf.reduce_sum(mask)
+
+    all_error = []
+    for batch in dataset:
+        batch = static_dataset.expand_batch(batch_to_dict(batch, signature_dynamic))
+        err = euclid_error(batch['mesh_mask'], batch['mesh_posed'], batch['mesh_body_posed'])
+        all_error.append(err.numpy())
+        # b = nest_conditions(static_dataset.expand_batch(batch_to_dict(d, signature_dynamic)))
+        print(len(all_error))
+
     
     dataset = DatasetSimple('/home/workspace/data/dataset/stylegan_new/train').batch(2)
 

docker/Dockerfile

@@ -1,31 +1,10 @@
-FROM tensorflow/tensorflow:2.3.1-gpu
+FROM tensorflow/tensorflow:2.9.0-gpu
-# FROM tensorflow/tensorflow:2.4.0rc0-gpu
-MAINTAINER moono.song "toilety@gmail.com"
-RUN apt-get update -y
-RUN apt-get install -y build-essential
-RUN pip install pillow
 
-# add sudoer
+RUN apt-get update
-ARG USERNAME=anonymous
+RUN apt-get upgrade -y
-ARG USER_UID=1000
+RUN apt-get install git ffmpeg libsm6 libxext6 libopenexr-dev zip unzip htop nano -y
-ARG USER_GID=$USER_UID
+RUN pip install numpy matplotlib opencv-python wandb OpenEXR tensorflow-graphics scipy plotly pandas tensorboardX opencv-python
 
-# Create the user
-RUN groupadd --gid $USER_GID $USERNAME \
-    && useradd --uid $USER_UID --gid $USER_GID -m $USERNAME \
-    #
-    # [Optional] Add sudo support. Omit if you don't need to install software after connecting.
-    && apt-get update \
-    && apt-get install -y sudo \
-    && echo $USERNAME ALL=\(root\) NOPASSWD:ALL > /etc/sudoers.d/$USERNAME \
-    && chmod 0440 /etc/sudoers.d/$USERNAME
 
-# ********************************************************
+WORKDIR /home/workspace/data
-# * Anything else you want to do like clean up goes here *
+# RUN wandb login
-# ********************************************************
-
-# [Optional] Set the default user. Omit if you want to keep the default as root.
-USER $USERNAME
-
-# add port for multiworker training
-EXPOSE 12345

docker/requierements.txt

@@ -6,3 +6,6 @@ OpenEXR
 tensorflow-graphics
 scipy
 plotly
+pandas
+tensorboardX
+opencv-python

load_models.py

@@ -18,8 +18,9 @@ def load_generator(g_params=None, is_g_clone=False, ckpt_dir=None, custom_cuda=T
         }
 
     test_latent = tf.ones((1, g_params['z_dim']), dtype=tf.float32)
-    test_conditions = (tf.ones((1, 72), dtype=tf.float32), tf.ones((1, 11), dtype=tf.float32), tf.ones((1, 3, 512, 256), dtype=tf.float32), tf.ones((1, 512, 256), dtype=tf.float32), tf.ones((1, 3, 512, 256), dtype=tf.float32))
+    # test_conditions = (tf.ones((1, 72), dtype=tf.float32), tf.ones((1, 11), dtype=tf.float32), tf.ones((1, 3, 512, 256), dtype=tf.float32), tf.ones((1, 512, 256), dtype=tf.float32), tf.ones((1, 3, 512, 256), dtype=tf.float32))
     # test_conditions = (tf.ones((1, 72), dtype=tf.float32), tf.ones((1, 11), dtype=tf.float32), tf.ones((1, 3, 512, 256), dtype=tf.float32))
+    test_conditions = { 'pose_shape': tf.ones((1, 83), dtype=tf.float32), 'uv_static': tf.ones((1, 3, 512, 256), dtype=tf.float32), 'uv_mask': tf.ones((1, 512, 256), dtype=tf.float32), 'uv_shaped_body': tf.ones((1, 3, 512, 256), dtype=tf.float32) }
 
     # build generator model
     generator = Generator(g_params)

losses.py

@@ -70,7 +70,10 @@ def g_logistic_ns_pathreg(real_images, generator, discriminator, z_dim,
 
     pl_minibatch = tf.maximum(1, tf.math.floordiv(batch_size, pl_minibatch_shrink))
     pl_z = tf.random.normal(shape=[pl_minibatch, z_dim], dtype=tf.float32)
-    pl_conditions = (conditions[0][:pl_minibatch], conditions[1][:pl_minibatch], conditions[2][:pl_minibatch], conditions[3][:pl_minibatch], conditions[4][:pl_minibatch])
+    pl_conditions = { key: conditions[key][:pl_minibatch] for key in conditions.keys() }
+    # for key in conditions.keys():
+    #     pl_conditions[key] = conditions[key][:pl_minibatch]
+    # pl_conditions = (conditions[0][:pl_minibatch], conditions[1][:pl_minibatch], conditions[2][:pl_minibatch], conditions[3][:pl_minibatch], conditions[4][:pl_minibatch])
 
     # forward pass
     fake_images = generator([z, conditions], ret_w_broadcasted=False, training=True)
@@ -124,7 +127,7 @@ def collision_loss(V, B, mask, thr=.004):
     D = V - tf.gather_nd(B, idx)
     b_vn = tf.gather_nd(B_vn, idx)
     dot = tf.einsum('abc,abc->ab', D, b_vn)
-    loss = tf.reduce_sum((tf.minimum(dot - thr, 0)  * mask) ** 2) / tf.reduce_sum(mask)
+    loss = tf.reduce_sum((tf.minimum(dot - thr, 0)  * mask) ** 2, axis=-1) # / tf.reduce_sum(mask, axis=-1)
 
     return loss
 
@@ -147,7 +150,7 @@ def edge_loss(x, y, mask):
     y_e = tf.sqrt(y_e + 1e-9)
     d_e = y_e - x_e
     d_e = d_e * e_mask
-    loss = tf.reduce_sum(d_e ** 2) / tf.reduce_sum(e_mask)
+    loss = tf.reduce_sum(d_e ** 2, axis=-1)# / tf.reduce_sum(e_mask, axis=-1)
     return loss
 
 

stylegan2/generator.py

@@ -105,8 +105,8 @@ class Generator(tf.keras.Model):
 
         image_out = self.synthesis(w_broadcasted)
         # image_out = image_out + conditions[2]
-        image_out = image_out + self.zero_body[:image_out.shape[0]]
+        image_out = image_out + conditions['uv_shaped_body']
-        image_out = image_out * conditions[3][:,tf.newaxis,:,:]
+        image_out = image_out * conditions['uv_mask'][:,tf.newaxis,:,:]
 
         if ret_w_broadcasted:
             return image_out, w_broadcasted

stylegan2/layers/embedding.py

@@ -97,6 +97,7 @@ class GarmentEmbedding(tf.keras.layers.Layer):
 
     def call(self, inputs, training=None, mask=None):
         x = inputs
+        x = tf.image.per_image_standardization(x)
 
         x = tf.transpose(x, [0,2,3,1])
         x = self.encoder(x, training=training)

stylegan2/layers/mapping_block.py

@@ -31,8 +31,8 @@ class Mapping(tf.keras.layers.Layer):
         x = latents
 
         # embed label if any
-        pose_shape_latent = self.pose_shape_embedding([conditions[0], conditions[1]])
+        pose_shape_latent = self.pose_shape_embedding([conditions['pose_shape'][:, :72], conditions['pose_shape'][:, 72:]])
-        garment_latent = self.garment_embedding(conditions[2])
+        garment_latent = self.garment_embedding(conditions['uv_static'])
         x = tf.concat([x, pose_shape_latent, garment_latent], axis=1)
         # x = tf.concat([x, garment_latent], axis=1)
 

train.py

@@ -5,20 +5,20 @@ import wandb
 from cv2 import drawFrameAxes
 import numpy as np
 import tensorflow as tf
-from custom_utils.custom_utils import UV2Mesh, maskFacesBatch, plot_3d
+from custom_utils.custom_utils import UV2Mesh, maskFacesBatch, nest_conditions, plot_3d, batch_to_dict
 from tensorboardX import SummaryWriter
 
 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 dataset import DatasetSimple, DatasetFactory, StaticDataset
 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
 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_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.')
+run = wandb.init(config=tf.compat.v1.flags.FLAGS, sync_tensorboard=True, project='stylegan', name='run_35', 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 and edge coeff. learns offsets over shaped body. fixed unposed uv map conversion. gnn discriminator. reworked data loader. per garment sum of errors. per_image_standardization for garment tempalte input')
 artifact = wandb.Artifact('code', type='code')
 artifact.add_file('train.py')
 artifact.add_file('losses.py')
@@ -40,7 +40,7 @@ smpl = SMPL()
 
 class Trainer(object):
     def __init__(self, t_params, name):
-        self.cur_tf_ver = t_params['cur_tf_ver']
+        # self.cur_tf_ver = t_params['cur_tf_ver']
         self.use_tf_function = t_params['use_tf_function']
         self.use_custom_cuda = t_params['use_custom_cuda']
         self.model_base_dir = t_params['model_base_dir']
@@ -129,12 +129,13 @@ class Trainer(object):
         return params_copy
 
     def d_train_step(self, dist_inputs):
-        uv_posed, conditions, _, _, _, _, _, uv_unposed = dist_inputs[0]
+        # uv_posed, conditions, _, _, _, _, _, uv_unposed = dist_inputs[0]
-        real_images = uv_unposed
+        batch = dist_inputs[0]
+        real_images = batch['uv_unposed']
 
         with tf.GradientTape() as d_tape:
             # compute losses
-            d_loss = d_logistic(real_images, self.generator, self.discriminator, self.g_params['z_dim'], conditions=conditions)
+            d_loss = d_logistic(real_images, self.generator, self.discriminator, self.g_params['z_dim'], conditions=batch['conditions'])
 
             # scale loss
             d_loss = tf.reduce_sum(d_loss) * self.global_batch_scaler
@@ -146,13 +147,14 @@ class Trainer(object):
         return d_loss
 
     def d_train_step_reg(self, dist_inputs):
-        uv_posed, conditions, _, _, _, _, _, uv_unposed = dist_inputs[0]
+        # uv_posed, conditions, _, _, _, _, _, uv_unposed = dist_inputs[0]
-        real_images = uv_unposed
+        batch = dist_inputs[0]
+        real_images = batch['uv_unposed']
 
         with tf.GradientTape() as d_tape:
             # compute losses
             d_gan_loss, r1_penalty = d_logistic_r1_reg(real_images, self.generator, self.discriminator,
-                                                       self.g_params['z_dim'], conditions=conditions)
+                                                       self.g_params['z_dim'], conditions=batch['conditions'])
             r1_penalty = r1_penalty * (0.5 * self.r1_gamma) * self.d_opt['reg_interval']
 
             # scale losses
@@ -168,12 +170,13 @@ class Trainer(object):
 
 
     def g_train_step(self, dist_inputs):
-        uv_posed, conditions, _, _, _, _, _, uv_unposed = dist_inputs[0]
+        # uv_posed, conditions, _, _, _, _, _, uv_unposed = dist_inputs[0]
-        real_images = uv_unposed
+        batch = dist_inputs[0]
+        real_images = batch['uv_unposed']
 
         with tf.GradientTape() as g_tape:
             # compute losses
-            _, g_loss = g_logistic_non_saturating(real_images, self.generator, self.discriminator, self.g_params['z_dim'], conditions=conditions)
+            _, g_loss = g_logistic_non_saturating(real_images, self.generator, self.discriminator, self.g_params['z_dim'], conditions=batch['conditions'])
 
             # scale loss
             g_loss = tf.reduce_sum(g_loss) * self.global_batch_scaler
@@ -183,14 +186,15 @@ class Trainer(object):
         return g_loss
 
     def g_train_step_reg(self, dist_inputs):
-        uv_posed, conditions, _, _, _, _, _, uv_unposed= dist_inputs[0]
+        # uv_posed, conditions, _, _, _, _, _, uv_unposed= dist_inputs[0]
-        real_images = uv_unposed
+        batch = dist_inputs[0]
+        real_images = batch['uv_unposed']
 
         with tf.GradientTape() as g_tape:
             # compute losses
             _, g_gan_loss, pl_penalty = g_logistic_ns_pathreg(real_images, self.generator, self.discriminator,
                                                            self.g_params['z_dim'], self.pl_mean,
-                                                           self.pl_minibatch_shrink, self.pl_denorm, self.pl_decay, conditions=conditions)
+                                                           self.pl_minibatch_shrink, self.pl_denorm, self.pl_decay, conditions=batch['conditions'])
             pl_penalty = pl_penalty * self.pl_weight * self.g_opt['reg_interval']
 
             # scale loss
@@ -206,28 +210,28 @@ class Trainer(object):
 
 
     def g_train_step_extra(self, dist_inputs):
-        uv_posed, conditions, mesh_static, mesh_body_posed, mesh_mask, mesh_posed, mesh_unposed, uv_unposed = dist_inputs[0]
+        # uv_posed, conditions, mesh_static, mesh_body_posed, mesh_mask, mesh_posed, mesh_unposed, uv_unposed = dist_inputs[0]
-        real_images = uv_unposed
+        batch = dist_inputs[0]
-        uv_mask = conditions[3]
+        real_images = batch['uv_unposed']
 
         with tf.GradientTape() as g_tape:
             # compute losses
-            fake_images, g_loss = g_logistic_non_saturating(real_images, self.generator, self.discriminator, self.g_params['z_dim'], conditions=conditions)
+            fake_images, g_loss = g_logistic_non_saturating(real_images, self.generator, self.discriminator, self.g_params['z_dim'], conditions=batch['conditions'])
 
             mesh_pred_unposed = uv2mesh(fake_images)
-            mesh_pred = smpl.run_garment(conditions[0], conditions[1], mesh_pred_unposed)
+            mesh_pred = smpl.run_garment(batch['conditions']['pose_shape'][:, :72], batch['conditions']['pose_shape'][:, 72:], mesh_pred_unposed)
 
-            e_loss = 1e4 * edge_loss(mesh_pred_unposed, mesh_static, mesh_mask)
+            e_loss = edge_loss(mesh_pred_unposed, batch['mesh_static'], batch['mesh_mask'])
-            c_loss = 1e2 * collision_penalty(mesh_pred, mesh_body_posed, mesh_mask)
+            c_loss = collision_penalty(mesh_pred, batch['mesh_body_posed'], batch['mesh_mask'])
-            l1_loss = 1e2 * tf.reduce_mean(tf.reduce_sum((tf.math.abs(real_images - fake_images)), axis=1) * uv_mask)
+            l1_loss = tf.reduce_sum(tf.reduce_sum((tf.math.abs(real_images - fake_images)), axis=1) * batch['conditions']['uv_mask'], axis=[1,2])
-            euclid_error = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(mesh_mask[:,:,tf.newaxis] * (mesh_posed - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(mesh_mask)
+            euclid_error = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (batch['mesh_posed'] - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
             ms_ssim = tf.image.ssim_multiscale(tf.transpose(real_images, [0,2,3,1]) + 1, tf.transpose(fake_images, [0,2,3,1]) + 1, max_val=2)
 
             # scale loss
             g_loss = tf.reduce_sum(g_loss) * self.global_batch_scaler
-            # e_loss = tf.reduce_sum(e_loss) * self.global_batch_scaler
+            e_loss = tf.reduce_sum(e_loss) * self.global_batch_scaler
-            # c_loss = tf.reduce_sum(c_loss) * self.global_batch_scaler
+            c_loss = tf.reduce_sum(c_loss) * self.global_batch_scaler
-            # l1_loss = tf.reduce_sum(l1_loss) * self.global_batch_scaler
+            l1_loss = tf.reduce_sum(l1_loss) * self.global_batch_scaler
             ms_ssim = tf.reduce_sum(ms_ssim) * self.global_batch_scaler
 
             g_loss = g_loss + self.unsupervised_coeff * (e_loss + c_loss) + self.gt_loss_coeff * l1_loss
@@ -238,32 +242,32 @@ class Trainer(object):
 
     def g_train_step_reg_extra(self, dist_inputs):
         # real_images, conditions, mesh_static, mesh_body_posed, mesh_mask, uv_mask = dist_inputs[0]
-        uv_posed, conditions, mesh_static, mesh_body_posed, mesh_mask, mesh_posed, mesh_unposed, uv_unposed = dist_inputs[0]
+        # uv_posed, conditions, mesh_static, mesh_body_posed, mesh_mask, mesh_posed, mesh_unposed, uv_unposed = dist_inputs[0]
-        real_images = uv_unposed
+        batch = dist_inputs[0]
-        uv_mask = conditions[3]
+        real_images = batch['uv_unposed']
 
         with tf.GradientTape() as g_tape:
             # compute losses
             fake_images, g_gan_loss, pl_penalty = g_logistic_ns_pathreg(real_images, self.generator, self.discriminator,
                                                            self.g_params['z_dim'], self.pl_mean,
-                                                           self.pl_minibatch_shrink, self.pl_denorm, self.pl_decay, conditions=conditions)
+                                                           self.pl_minibatch_shrink, self.pl_denorm, self.pl_decay, conditions=batch['conditions'])
             pl_penalty = pl_penalty * self.pl_weight * self.g_opt['reg_interval']
 
             mesh_pred_unposed = uv2mesh(fake_images)
-            mesh_pred = smpl.run_garment(conditions[0], conditions[1], mesh_pred_unposed)
+            mesh_pred = smpl.run_garment(batch['conditions']['pose_shape'][:, :72], batch['conditions']['pose_shape'][:, 72:], mesh_pred_unposed)
 
-            e_loss = 1e4 * edge_loss(mesh_pred_unposed, mesh_static, mesh_mask)
+            e_loss = edge_loss(mesh_pred_unposed, batch['mesh_static'], batch['mesh_mask'])
-            c_loss = 1e2 * collision_penalty(mesh_pred, mesh_body_posed, mesh_mask)
+            c_loss = collision_penalty(mesh_pred, batch['mesh_body_posed'], batch['mesh_mask'])
-            l1_loss = 1e2 * tf.reduce_sum(tf.reduce_sum((tf.math.abs(real_images - fake_images)), axis=1) * uv_mask) / tf.reduce_sum(uv_mask)
+            l1_loss = tf.reduce_sum(tf.reduce_sum((tf.math.abs(real_images - fake_images)), axis=1) * batch['conditions']['uv_mask'], axis=[1,2])
-            euclid_error = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(mesh_mask[:,:,tf.newaxis] * (mesh_posed - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(mesh_mask)
+            euclid_error = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (batch['mesh_posed'] - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
             ms_ssim = tf.image.ssim_multiscale(tf.transpose(real_images, [0,2,3,1]) + 1, tf.transpose(fake_images, [0,2,3,1]) + 1, max_val=2)
 
             # scale loss
             pl_penalty = tf.reduce_sum(pl_penalty) * self.global_batch_scaler
             g_gan_loss = tf.reduce_sum(g_gan_loss) * self.global_batch_scaler
-            # e_loss = tf.reduce_sum(e_loss) * self.global_batch_scaler
+            e_loss = tf.reduce_sum(e_loss) * self.global_batch_scaler
-            # c_loss = tf.reduce_sum(c_loss) * self.global_batch_scaler
+            c_loss = tf.reduce_sum(c_loss) * self.global_batch_scaler
-            # l1_loss = tf.reduce_sum(l1_loss) * self.global_batch_scaler
+            l1_loss = tf.reduce_sum(l1_loss) * self.global_batch_scaler
             ms_ssim = tf.reduce_sum(ms_ssim) * self.global_batch_scaler
             
             # combine
@@ -273,7 +277,7 @@ class Trainer(object):
         self.g_optimizer.apply_gradients(zip(g_gradients, self.generator.trainable_variables))
         return g_loss, g_gan_loss, pl_penalty, e_loss, c_loss, l1_loss, euclid_error, ms_ssim
 
-    def train(self, dist_datasets, strategy):
+    def train(self, dist_datasets: tf.data.Dataset, static_dataset: StaticDataset, dataset_signature: dict, strategy):
         @tf.function
         def dist_d_train_step(inputs):
             per_replica_losses = strategy.run(fn=self.d_train_step, args=(inputs,))
@@ -306,9 +310,9 @@ class Trainer(object):
         def dist_g_train_step_extra(inputs):
             per_replica_losses = strategy.run(fn=self.g_train_step_extra, args=(inputs,))
             mean_g_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[0], axis=None)
-            mean_e_loss = strategy.reduce(tf.distribute.ReduceOp.MEAN, per_replica_losses[1], axis=None)
+            mean_e_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[1], axis=None)
-            mean_c_loss = strategy.reduce(tf.distribute.ReduceOp.MEAN, per_replica_losses[2], axis=None)
+            mean_c_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[2], axis=None)
-            mean_l1_loss = strategy.reduce(tf.distribute.ReduceOp.MEAN, per_replica_losses[3], axis=None)
+            mean_l1_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[3], axis=None)
             mean_euclid_error = strategy.reduce(tf.distribute.ReduceOp.MEAN, per_replica_losses[4], axis=None)
             mean_ms_ssim = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[5], axis=None)
             return mean_g_loss, mean_e_loss, mean_c_loss, mean_l1_loss, mean_euclid_error, mean_ms_ssim
@@ -319,9 +323,9 @@ class Trainer(object):
             mean_g_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[0], axis=None)
             mean_g_gan_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[1], axis=None)
             mean_pl_penalty = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[2], axis=None)
-            mean_e_loss = strategy.reduce(tf.distribute.ReduceOp.MEAN, per_replica_losses[3], axis=None)
+            mean_e_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[3], axis=None)
-            mean_c_loss = strategy.reduce(tf.distribute.ReduceOp.MEAN, per_replica_losses[4], axis=None)
+            mean_c_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[4], axis=None)
-            mean_l1_loss = strategy.reduce(tf.distribute.ReduceOp.MEAN, per_replica_losses[5], axis=None)
+            mean_l1_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[5], axis=None)
             mean_euclid_error = strategy.reduce(tf.distribute.ReduceOp.MEAN, per_replica_losses[6], axis=None)
             mean_ms_ssim = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[7], axis=None)
             return mean_g_loss, mean_g_gan_loss, mean_pl_penalty, mean_e_loss, mean_c_loss, mean_l1_loss, mean_euclid_error, mean_ms_ssim
@@ -361,6 +365,7 @@ class Trainer(object):
         t_start = time.time()
         step = self.g_optimizer.iterations.numpy()
         for batch in dist_datasets:
+            batch = nest_conditions(static_dataset.expand_batch(batch_to_dict(batch, dataset_signature)))
 
             # d train step
             if (step + 1) % self.d_opt['reg_interval'] == 0:
@@ -420,9 +425,6 @@ class Trainer(object):
             metric_euclid_error(euclid_error)
             metric_ms_ssim(ms_ssim)
 
-            # get current step
-            step += 1
-
             # save to tensorboard
             with train_summary_writer.as_default():
                 tf.summary.scalar('d_loss', metric_d_loss.result(), step=step)
@@ -447,21 +449,18 @@ class Trainer(object):
             if step % self.image_summary_step == 0:
                 # add summary image
                 test_z = tf.random.normal(shape=(self.n_samples, self.g_params['z_dim']), dtype=tf.dtypes.float32)
-                test_conditions = batch[1] #tf.ones((self.n_samples, self.g_params['labels_dim']), dtype=tf.dtypes.float32)
+                test_conditions = batch['conditions'] #tf.ones((self.n_samples, self.g_params['labels_dim']), dtype=tf.dtypes.float32)
                 summary_image = dist_gen_samples((test_z, test_conditions))
 
                 mesh_unposed = uv2mesh(summary_image[:, :, :self.train_res, int(self.train_res/2):])
-                mesh = smpl.run_garment(batch[1][0], batch[1][1], mesh_unposed)
+                mesh = smpl.run_garment(batch['conditions']['pose_shape'][:, :72], batch['conditions']['pose_shape'][:, 72:], mesh_unposed)
                 # convert to tensor image
                 summary_image = self.convert_per_replica_image(summary_image, strategy)
 
                 with train_summary_writer.as_default():
                     tf.summary.image('images', summary_image, step=step)
 
-                template = uv2mesh(batch[1][2])
+                plot_meshes(mesh_writer, garment_unposed=mesh_unposed, garment=mesh, mesh_gt=batch['mesh_unposed'], template=batch['mesh_static'], body=batch['mesh_body_posed'], mask=batch['mesh_mask'].numpy(), step=step)
-                # gt = uv2mesh(batch[0])
-                plot_meshes(mesh_writer, garment_unposed=mesh_unposed, garment=mesh, mesh_gt=batch[6], template=template, body=batch[3], mask=batch[4].numpy(), step=step)
-                # plot_meshes(mesh_writer, garment=gt, template=template, body=batch[3], mask=batch[4].numpy(), step=step)
 
                 # train_summary_writer.flush()
 
@@ -479,6 +478,9 @@ class Trainer(object):
             if step >= self.max_steps:
                 break
 
+            # get current step
+            step += 1
+
         # save last checkpoint
         step = self.g_optimizer.iterations.numpy()
         self.manager.save(checkpoint_number=step)
@@ -492,7 +494,7 @@ class Trainer(object):
         fake_images_07 = self.g_clone([test_z, test_conditions], truncation_psi=0.7, training=False)
 
         # merge on batch dimension: [n_samples, 3, out_res, 2 * out_res]
-        template = tf.concat([test_conditions[2], test_conditions[2]], axis=2)
+        template = tf.concat([test_conditions['uv_static'], test_conditions['uv_static']], axis=2)
         fake_image = tf.concat([fake_images_05, fake_images_07], axis=2)
         final_image = tf.concat([template, fake_image], axis=3)
 
@@ -552,10 +554,11 @@ def main():
     parser.add_argument('--train_res', default=512, type=int)
     parser.add_argument('--shuffle_buffer_size', default=1000, type=int)
     parser.add_argument('--batch_size_per_replica', default=4, type=int)
+    parser.add_argument('--name', type=str)
     args = vars(parser.parse_args())
 
     # check tensorflow version
-    cur_tf_ver = check_tf_version()
+    # cur_tf_ver = check_tf_version()
 
     # GPU environment settings
     if args['allow_memory_growth']:
@@ -594,7 +597,29 @@ 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, drop_remainder=True).repeat()
+    # dataset = DatasetSimple('/home/workspace/data/dataset/stylegan_new/train').cache().batch(global_batch_size, drop_remainder=True).repeat()
+
+    factory = DatasetFactory()
+    signature_dynamic = { 'id': (),
+                          'pose_shape': (72 + 11,),
+                        #   'uv_posed': (3, 512, 256,),
+                        #   'uv_body_posed': (3, 512, 256,),
+                          'mesh_body_posed': (14475, 3,),
+                          'mesh_posed': (14475, 3,),
+                          'mesh_unposed': (14475, 3,),
+                          'uv_unposed': (3, 512, 256,) }
+    signature_static = { 'uv_static': (3, 512, 256,),
+                         'mesh_static': (14475, 3,),
+                         'mesh_mask': (14475,),
+                         'uv_mask': (512, 256,),
+                         'uv_shaped_body': (3, 512, 256,), }
+
+    dataset = factory.get_dataset(folder='/home/workspace/data/dataset/stylegan_new/train',
+                                  batch_size=4,
+                                  signature=signature_dynamic)
+    static_dataset = StaticDataset(folder='/home/workspace/data/dataset/stylegan_new/train',
+                                   signature=signature_static)
+
 
     options = tf.data.Options()
     options.experimental_distribute.auto_shard_policy = tf.data.experimental.AutoShardPolicy.OFF
@@ -607,7 +632,7 @@ def main():
         # training parameters
         training_parameters = {
             # global params
-            'cur_tf_ver': cur_tf_ver,
+            # 'cur_tf_ver': cur_tf_ver,
             'use_tf_function': args['use_tf_function'],
             'use_custom_cuda': args['use_custom_cuda'],
             'model_base_dir': args['model_base_dir'],
@@ -627,8 +652,8 @@ def main():
             'train_res': args['train_res'],
         }
 
-        trainer = Trainer(training_parameters, name=f'stylegan2-test-33-{args["train_res"]}x{args["train_res"]}')
+        trainer = Trainer(training_parameters, name=f'stylegan2-test-35-{args["train_res"]}x{args["train_res"]}')
-        trainer.train(dist_dataset, strategy)
+        trainer.train(dist_dataset, static_dataset, signature_dynamic, strategy)
     return