lot of things

.gitignore

@@ -142,4 +142,5 @@ models/
 models_archive/
 temp.jpg
 corrupt/
-wandb copy/
+wandb copy/
+checkpoints/

auto_encoder.py

@@ -1,7 +1,6 @@
 import os
 import time
 import argparse
-from cv2 import polarToCart
 import wandb
 import random
 import numpy as np
@@ -9,6 +8,7 @@ import tensorflow as tf
 from custom_utils.custom_utils import UV2Mesh, maskFacesBatch
 from tensorboardX import SummaryWriter
 import plotly.graph_objects as go
+import plotly.express as ex
 
 from utils import str_to_bool
 from tf_utils import check_tf_version, allow_memory_growth, split_gpu_for_testing
@@ -16,16 +16,25 @@ from load_models import load_generator, load_discriminator
 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 convnext.models.convnext_tf import ConvNeXt
+from convnext.models.convnext_tf import ConvNeXt, ConvNeXtReverse, Block
 from stylegan2_ae.generator import Generator
+from custom_utils.graph_utils import Reconstructor_V4_Decoder as GNNDecoder
+from custom_utils.graph_utils import Reconstructor_v2 as SimpleReconstructor
+from dataset import DatasetFactory, StaticDataset
+from custom_utils.custom_utils import batch_to_dict
+from custom_utils.smpl import SMPL
 
-
+name='ae_body_1'
-os.environ["CUDA_VISIBLE_DEVICES"] = "2"
+# os.environ["CUDA_VISIBLE_DEVICES"] = "6"
-run = wandb.init(config=tf.compat.v1.flags.FLAGS, sync_tensorboard=True, project='stylegan', name='ae_unposed_1', notes='convnext + stylegan generator')
+# run = wandb.init(project='stylegan', name=name,
-artifact = wandb.Artifact('code', type='code')
+#          notes='convnext + convnext, add dense with tanh at the end of reverse convnext, larger features. add lr scheduler. fft3d instead of rfft3d. \
-artifact.add_file('auto_encoder.py')
+# fft2d per channel and separate for top and bottom. fft_loss, all data. added convnext block as bottleneck, same as before, added extra metric. \
-artifact.add_file('stylegan2_ae/layers/mapping_block.py')
+# . lowered learning rate because it goes to nan. shallow third stage. slower lr decay. with validation. higher lr rate.')
-run.log_artifact(artifact)
+# artifact = wandb.Artifact('code', type='code')
+# artifact.add_file('auto_encoder.py')
+# artifact.add_file('stylegan2_ae/layers/mapping_block.py')
+# artifact.add_file('convnext/models/convnext_tf.py')
+# run.log_artifact(artifact)
 
 uv2mesh = UV2Mesh(res=256)
 
@@ -67,8 +76,46 @@ def filter_resolutions_featuremaps(resolutions, featuremaps, res):
     return filtered_resolutions, filtered_featuremaps
 
 
-def plot_3d(mesh_gt, mesh_pred, faces):
+# def plot_3d(mesh_gt, mesh_pred, faces):
-    mesh_pred[:,1] += 2
+#     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
+
+
+def plot_3d(mesh_gt, mesh_pred_unposed, faces, collision, collision_faces, collision_color):
+    mesh_pred_unposed[:,1] += 1
+    collision[:,1] += 2
     fig = go.Figure(data=[
         go.Mesh3d(
             # 8 vertices of a cube
@@ -82,13 +129,24 @@ def plot_3d(mesh_gt, mesh_pred, faces):
         ),
         go.Mesh3d(
             # 8 vertices of a cube
-            x=mesh_pred[:,0],
+            x=mesh_pred_unposed[:,0],
-            y=mesh_pred[:,1],
+            y=mesh_pred_unposed[:,1],
-            z=mesh_pred[:,2],
+            z=mesh_pred_unposed[:,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=collision[:,0],
+            y=collision[:,1],
+            z=collision[:,2],
+            # i, j and k give the vertices of triangles
+            i = collision_faces[:,0],
+            j = collision_faces[:,1],
+            k = collision_faces[:,2],
+            vertexcolor=collision_color
         )
     ])
 
@@ -103,31 +161,35 @@ def plot_3d(mesh_gt, mesh_pred, faces):
 
     return fig
 
+class Bottleneck(tf.keras.Model):
+    
+    def __init__(self, dim):
+        super(Bottleneck, self).__init__()
+
+        self.bottleneck = tf.keras.Sequential(
+                [Block(dim=dim, drop_path=0,
+                       layer_scale_init_value=1e-6, prefix=f'bottleneck{j}') for j in range(5)]
+            )
+
+    def call(self, inputs, training=None):
+        x = inputs
+
+        x = self.bottleneck(x, training=training)
+        
+        return x
+
 
 class AutoEncoder(tf.keras.Model):
     
     def __init__(self):
         super(AutoEncoder, self).__init__()
 
-        resolutions = [4, 8, 16, 32, 64, 128, 256, 512, 1024]
-        g_featuremaps = [256, 256, 256, 256, 256, 128, 64, 32, 16]
-        train_resolutions, train_g_featuremaps = filter_resolutions_featuremaps(resolutions, g_featuremaps, 512)
-        g_params = {
-            'z_dim': 256,
-            'w_dim': 256,
-            'labels_dim': -1,
-            'n_mapping': 8,
-            'resolutions': train_resolutions,
-            'featuremaps': train_g_featuremaps,
-        }
 
-        # test_latent = tf.ones((1, g_params['z_dim']), dtype=tf.float32)
+        bottleneck_size = 32
-        self.decoder = Generator(g_params)
-        # _ = generator(test_latent)
 
-
+        self.encoder = ConvNeXt(depths=[3, 3, 9, 3], dims=[32, 64, 128, bottleneck_size], include_top=False)
-        self.encoder = ConvNeXt(depths=[3, 3, 27, 3], dims=[96, 192, 256, g_params['z_dim']], include_top=False)
+        self.bottleneck = Bottleneck(bottleneck_size)
-        self.pooling = tf.keras.layers.GlobalAveragePooling2D()
+        self.decoder = ConvNeXtReverse(depths=list(reversed([3, 3, 9, 3])), dims=list(reversed([bottleneck_size, 128, 64, 16])))
 
 
     def call(self, inputs, training=None):
@@ -135,105 +197,237 @@ class AutoEncoder(tf.keras.Model):
 
         x = tf.transpose(x, [0,2,3,1])
         x = self.encoder(x, training=training)
-        x = self.pooling(x)
+        # x = self.pooling(x)
+        # x = self.reshape(x)
+        # x = self.flatten(x)
+
+        x = self.bottleneck(x, training=training)
 
         x = self.decoder(x, training=training)
+        x = tf.transpose(x, [0,3,1,2])
+
+        # x = x * uv_mask[:,:, tf.newaxis]
         x = x * uv_mask[:,tf.newaxis,:,:]
 
         return x
 
+def draw_cicle(shape,diamiter):
+    '''
+    Input:
+    shape    : tuple (height, width)
+    diameter : scalar
+    
+    Output:
+    np.array of shape  that says True within a circle with diamiter =  around center 
+    '''
+    assert len(shape) == 2
+    TF = np.zeros(shape,dtype=bool)
+    center = np.array(TF.shape)/2.0
+
+    for iy in range(shape[0]):
+        for ix in range(shape[1]):
+            TF[iy,ix] = (iy- center[0])**2 + (ix - center[1])**2 < diamiter **2
+    return (TF)
+
+
+TFcircleIN   = draw_cicle(shape=(256,256), diamiter=15)
+TFcircleIN = np.ones((512,256,3), dtype=bool) * np.concatenate((TFcircleIN, TFcircleIN), axis=0)[:,:,np.newaxis]
+TFcircleOUT  = ~TFcircleIN
+
+def fft_loss_per_channel(pred, gt):
+    pred = tf.complex(pred, tf.zeros_like(pred))
+    gt = tf.complex(gt, tf.zeros_like(gt))
+
+    fft_pred_top = tf.stack([tf.signal.fftshift(tf.signal.fft2d(pred[:,i,:256,:])) for i in range(3)], axis=-1)
+    fft_gt_top = tf.stack([tf.signal.fftshift(tf.signal.fft2d(gt[:,i,:256,:])) for i in range(3)], axis=-1)
+    
+    fft_pred_bottom = tf.stack([tf.signal.fftshift(tf.signal.fft2d(pred[:,i,256:,:])) for i in range(3)], axis=-1)
+    fft_gt_bottom = tf.stack([tf.signal.fftshift(tf.signal.fft2d(gt[:,i,256:,:])) for i in range(3)], axis=-1)
+    
+    fft_pred = tf.concat((fft_pred_top, fft_pred_bottom), axis=1)
+    fft_gt = tf.concat((fft_gt_top, fft_gt_bottom), axis=1)
+
+    fft_pred_low = tf.where(TFcircleIN, fft_pred, tf.zeros_like(fft_pred))
+    fft_gt_low = tf.where(TFcircleIN, fft_gt, tf.zeros_like(fft_gt))
+    fft_pred_high = tf.where(TFcircleOUT, fft_pred, tf.zeros_like(fft_pred))
+    fft_gt_high = tf.where(TFcircleOUT, fft_gt, tf.zeros_like(fft_gt))
+
+    low_loss = tf.reduce_mean(tf.reduce_sum(tf.math.abs(fft_pred_low - fft_gt_low), axis=[1,2,3]))
+    high_loss = tf.reduce_mean(tf.reduce_sum(tf.math.abs(fft_pred_high - fft_gt_high), axis=[1,2,3]))
+    total = tf.reduce_mean(tf.reduce_sum(tf.math.abs(fft_pred - fft_gt), axis=[1,2,3]))
+
+    return total, low_loss, high_loss
+
+    # ex.imshow(tf.math.real(gt[0,0,:256,:])).show()
+    # img = tf.signal.fftshift(fft_gt[0][0])
+    # ex.imshow(tf.math.abs(tf.math.log(img))).show()
+    # low = tf.where(TFcircleIN, img, tf.zeros_like(img))
+    # ex.imshow(tf.math.real(tf.signal.ifft2d(tf.signal.ifftshift(low)))).show()
+    # high = tf.where(TFcircleOUT, img, tf.zeros_like(img))
+    # # ex.imshow(tf.math.abs(tf.math.log(high))).show()
+    # ex.imshow(tf.math.real(tf.signal.ifft2d(tf.signal.ifftshift(high)))).show()
+
+
+
+
+    print()
+
 
 @tf.function
 def train_step(batch):
-    uv_unposed, mesh_unposed, uv_mask, mesh_mask = batch
 
     with tf.GradientTape() as tape:
 
-        pred = model([uv_unposed, uv_mask])
+        pred = model([batch['uv_unposed'], batch['uv_mask']])
+        # pred = (pred + zero_body) * batch['mesh_mask'][:,:,tf.newaxis]
 
         mesh_pred = uv2mesh(pred)
-        mesh_orig = mesh_unposed
+        # mesh_pred = pred
+        mesh_orig = batch['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(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)
+        e_loss = 1e-2 * tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']))
+        c_loss = 1e-2 * tf.reduce_mean(collision_penalty(mesh_pred, batch['mesh_body_posed'], batch['mesh_mask']))
+        # l1_loss = tf.reduce_mean(tf.reduce_sum((tf.math.abs(batch['mesh_unposed'] - pred)), axis=-1) * batch['mesh_mask'])
+        l1_loss = tf.reduce_mean(tf.reduce_sum((tf.math.abs(batch['uv_unposed'] - pred)), axis=1) * batch['uv_mask'])
+        ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(tf.transpose(pred, [0,2,3,1]) + 1, tf.transpose(batch['uv_unposed'], [0,2,3,1]) + 1, max_val=2))
+        rfft_loss = 0 #tf.reduce_mean(tf.reduce_sum(tf.math.abs(tf.signal.rfft3d(pred) - tf.signal.rfft3d(batch['uv_unposed'])), axis=[1,2,3]))
+        fft_loss = tf.reduce_mean(tf.reduce_sum(tf.math.abs(tf.signal.fft3d(tf.complex(pred, tf.zeros_like(pred))) - tf.signal.fft3d(tf.complex(batch['uv_unposed'], tf.zeros_like(batch['uv_unposed'])))), axis=[1,2,3]))
+        fft_loss_pc, fft_loss_pc_low, fft_loss_pc_high = 0,0,0 #fft_loss_per_channel(pred, batch['uv_unposed'])
+        # loss = fft_loss_pc # + e_loss + c_loss
+        loss = 1e3 * (1 - ms_ssim)
+
+    mesh_orig_2 = uv2mesh(batch['uv_unposed'])
+    # mesh_skinned = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred)
+    metric = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
+    metric_2 = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig_2 - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
+    euclid_posed = 0 # tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_skinned - batch['mesh_posed']) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
     gradients = tape.gradient(loss, model.trainable_variables)
+    # gradients = [tf.clip_by_norm(g, 1.0) for g in gradients]
+    # gradients = [tf.clip_by_value(g, -1.0, 1.0) for g in gradients]
     optimizer.apply_gradients(zip(gradients, model.trainable_variables))
-    return { 'loss': loss, 'edge_loss': e_loss, 'collision_loss': c_loss, 'l1_loss': l1_loss, 'euclid_error': metric }
+    return { 'train/loss': loss, 'train/edge_loss': e_loss, 'train/collision_loss': c_loss, 'train/l1_loss': l1_loss, 'train/euclid_error': metric, 'train/ms_ssim': ms_ssim,
+             'train/fft_loss': fft_loss, 'train/euclid_error_reconst': metric_2, 'train/rfft_loss': rfft_loss, 'train/fft_oc_low': fft_loss_pc_low, 'train/fft_pc_high': fft_loss_pc_high,
+             'train/fft_loss_per_channel': fft_loss_pc, 'train/euclid_error_posed': euclid_posed }
         
 
 
 @tf.function
 def valid_step(batch):
-    uv_unposed, mesh_unposed, uv_mask, mesh_mask = batch
+    pred = model([batch['uv_unposed'], batch['uv_mask']])
 
-    pred = model([uv_unposed, uv_mask])
     mesh_pred = uv2mesh(pred)
-    mesh_orig = mesh_unposed
+    mesh_orig = batch['mesh_unposed']
+    e_loss = 1e-2 * tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']))
+    c_loss = 1e-2 * tf.reduce_mean(collision_penalty(mesh_pred, batch['mesh_body_posed'], batch['mesh_mask']))
+    l1_loss = tf.reduce_mean(tf.reduce_sum((tf.math.abs(batch['uv_unposed'] - pred)), axis=1) * batch['uv_mask'])
+    ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(tf.transpose(pred, [0,2,3,1]) + 1, tf.transpose(batch['uv_unposed'], [0,2,3,1]) + 1, max_val=2))
+    fft_loss = tf.reduce_mean(tf.reduce_sum(tf.math.abs(tf.signal.fft3d(tf.complex(pred, tf.zeros_like(pred))) - tf.signal.fft3d(tf.complex(batch['uv_unposed'], tf.zeros_like(batch['uv_unposed'])))), axis=[1,2,3]))
+    loss = 1e3 * (1 - ms_ssim)
+    mesh_orig_2 = uv2mesh(batch['uv_unposed'])
+    metric = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
+    metric_2 = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig_2 - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
+  
+    return { 'valid/loss': loss, 'valid/edge_loss': e_loss, 'valid/collision_loss': c_loss, 'valid/l1_loss': l1_loss, 'valid/euclid_error': metric, 'valid/ms_ssim': ms_ssim,
+             'valid/fft_loss': fft_loss, 'valid/euclid_error_reconst': metric_2
+            }
 
-    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(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)
+def plot_everything(batch: dict, type: str):
-    return { 'loss': loss, 'edge_loss': e_loss, 'collision_loss': c_loss, 'l1_loss': l1_loss, 'euclid_error': metric }
+
+    pred = model([batch['uv_unposed'], batch['uv_mask']])
+    
+    mesh_pred = uv2mesh(pred)
+    mesh_orig = batch['mesh_unposed']
+    
+    orig_img = tf.transpose(batch['uv_unposed'], [0,2,3,1])
+    pred_img = tf.transpose(pred, [0,2,3,1])
+    img = tf.concat((orig_img, pred_img), axis=2)
+    imgs = [wandb.Image((img[i] + 1) / 2.0) for i in range(4)]
+
+    F = np.load("custom_utils/smpl_faces.npy")
+    mask = batch['mesh_mask'].numpy().astype(bool)
+    garment_unposed, faces = maskFacesBatch(mesh_pred.numpy(), F, mask)
+    gt = [mesh_orig[i].numpy()[mask[i], :] for i in range(pred.shape[0])]
+    
+    body = batch['mesh_body_posed'].numpy()
+    mesh_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred).numpy()
+    garment_posed = [mesh_posed[i, mask[i], :] for i in range(body.shape[0])]
+    collision = [np.concatenate((body[i], garment_posed[i]), axis=0) for i in range(len(garment_posed))]
+    collision_faces = [np.concatenate((F, faces[i] + 14475), axis=0) for i in range(len(garment_posed))]
+    collision_color = [np.concatenate((np.ones_like(body[i], dtype=np.int32)*100, np.ones_like(garment_posed[i], dtype=np.int32)*230), axis=0) for i in range(len(garment_posed))]
+
+    figs = { type + "/garment_" + str(i): plot_3d(gt[i], garment_unposed[i], faces[i], collision[i], collision_faces[i], collision_color[i]) for i in range(4) }
+    
+    return { **figs, **{ type + '/images': imgs } }
+
 
 
 if __name__ == '__main__':
     allow_memory_growth()
+    parser = argparse.ArgumentParser(description='')
+    parser.add_argument('--name', type=str)
+    args = vars(parser.parse_args())
     
-    
+    smpl = SMPL()
-    name = 'ae_unposed_1'
+    # R = smpl.run_rodrigues((np.zeros((8,72), dtype=np.float32)))
-    dataset_train = DatasetSimple('/home/workspace/data/dataset/stylegan_new/train').cache().batch(8).repeat()
+    # zero_body = smpl.run(R, np.zeros((8, 11), dtype=np.float32))[0]
-    dataset_valid = DatasetSimple('/home/workspace/data/dataset/stylegan_new/valid').cache().batch(8).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_train = factory.get_dataset(folder='/home/workspace/local-data/jumpsuit_20/train',
+                                  batch_size=8,
+                                  signature=signature_dynamic)
+    static_dataset_train = StaticDataset(folder='/home/workspace/local-data/jumpsuit_20/train',
+                                   signature=signature_static)
+    dataset_valid = factory.get_dataset(folder='/home/workspace/local-data/jumpsuit_20/valid',
+                                  batch_size=8,
+                                  signature=signature_dynamic)
+    static_dataset_valid = StaticDataset(folder='/home/workspace/local-data/jumpsuit_20/valid',
+                                   signature=signature_static)
+
 
     model = AutoEncoder()
     # optimizer = tf.optimizers.Adam(learning_rate=0.002, beta_1=0.0, beta_2=0.99, epsilon=1e-8)
-    optimizer = tf.optimizers.Adam()
+    lr_schedule = tf.keras.optimizers.schedules.ExponentialDecay(0.0005, decay_steps=30000, decay_rate=0.1, staircase=False)
+    optimizer = tf.optimizers.Adam(learning_rate=lr_schedule)
     folder = 'models/ae_logs/'
-    train_summary_writer = tf.summary.create_file_writer(folder + name + '/train')
+    # train_summary_writer = tf.summary.create_file_writer(folder + name + '/train')
     # valid_summary_writer = tf.summary.create_file_writer(folder + name + '/valid')
+    valid_iterator = iter(dataset_valid)
 
     step = 0
     for batch in dataset_train:
+        batch = static_dataset_train.expand_batch(batch_to_dict(batch, signature_dynamic))
         train_loss = train_step(batch)
 
-        for key in train_loss.keys():
+        data_valid = valid_iterator.get_next()
-            with train_summary_writer.as_default():
+        data_valid = static_dataset_valid.expand_batch(batch_to_dict(data_valid, signature_dynamic))
-                tf.summary.scalar(key, train_loss[key], step=step)
+        valid_loss = valid_step(data_valid)
 
-        data_valid = batch
+        if step % 1000 == 0:
-        # data_valid = iter(dataset_valid).get_next()
+            train_figs = plot_everything(batch, 'train')
-        # valid_loss = valid_step(data_valid)
+            valid_figs = plot_everything(data_valid, 'valid')
-
+            wandb.log({ **train_loss, **valid_loss, **train_figs, **valid_figs }, step=step)
-        # for key in valid_loss.keys():
+        else:
-        #     with valid_summary_writer.as_default():
+            wandb.log({ **train_loss, **valid_loss }, step=step)
-        #         tf.summary.scalar(key, valid_loss[key], step=step)
-
-        if step % 250 == 0:
-            pred = model([data_valid[0], data_valid[2]])
             
-            mesh_pred = uv2mesh(pred)
-            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])
-            img = tf.concat((orig_img, pred_img), axis=2)
-            with train_summary_writer.as_default():
-                tf.summary.image('images', img, step=step)
-
-            F = np.load("custom_utils/smpl_faces.npy")
-            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(pred.shape[0]):
-                fig = plot_3d(gt[i], garment[i], faces[i])
-                wandb.log({"garment_" + str(i): fig})
-                
 
         print(step)
 

auto_encoder_body.py

@@ -0,0 +1,435 @@
+import os
+import time
+import argparse
+import wandb
+import random
+import cv2
+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
+import plotly.express as ex
+
+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 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 convnext.models.convnext_tf import ConvNeXt, ConvNeXtReverse, Block
+from stylegan2_ae.generator import Generator
+from custom_utils.graph_utils import Reconstructor_V4_Decoder as GNNDecoder
+from custom_utils.graph_utils import Reconstructor_v2 as SimpleReconstructor
+from dataset import DatasetFactory, StaticDataset
+from custom_utils.custom_utils import batch_to_dict
+from custom_utils.smpl import SMPL
+
+name='ae_body_5'
+def init_wandb():
+    os.environ["CUDA_VISIBLE_DEVICES"] = "3"
+    run = wandb.init(project='vae', name=name,
+            notes='convnext + convnext, add dense with tanh at the end of reverse convnext, larger features. add lr scheduler. fft3d instead of rfft3d. \
+    fft2d per channel and separate for top and bottom. fft_loss, all data. added convnext block as bottleneck, same as before, added extra metric. \
+    . lowered learning rate because it goes to nan. shallow third stage. slower lr decay. with validation. higher lr rate.\
+    two more downscaling than up, so decoder output is resized wto final resolution\
+    larger batch_size.\
+    flattened bottleneck with 128 latent_dim\ remove global orientation')
+    artifact = wandb.Artifact('code', type='code')
+    artifact.add_file('auto_encoder_body.py')
+    artifact.add_file('losses.py')
+    artifact.add_file('dataset.py')
+    artifact.add_file('convnext/models/convnext_tf.py')
+    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)
+    filtered_resolutions = resolutions[:index + 1]
+    filtered_featuremaps = featuremaps[:index + 1]
+    return filtered_resolutions, filtered_featuremaps
+
+
+def plot_3d(mesh_gt, mesh_pred_unposed, faces, collision, collision_faces, collision_color):
+    mesh_pred_unposed[:,1] += 1
+    collision[:,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_unposed[:,0],
+            y=mesh_pred_unposed[:,1],
+            z=mesh_pred_unposed[:,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=collision[:,0],
+            y=collision[:,1],
+            z=collision[:,2],
+            # i, j and k give the vertices of triangles
+            i = collision_faces[:,0],
+            j = collision_faces[:,1],
+            k = collision_faces[:,2],
+            vertexcolor=collision_color
+        )
+    ])
+
+    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 Bottleneck(tf.keras.Model):
+    
+    def __init__(self, dim, latent_dim):
+        super(Bottleneck, self).__init__()
+
+        self.bottleneck = tf.keras.Sequential(
+                [Block(dim=dim, drop_path=0,
+                       layer_scale_init_value=1e-6, prefix=f'bottleneck{j}') for j in range(5)]
+            )
+        self.flatten = tf.keras.layers.Flatten()
+        self.dense = tf.keras.layers.Dense(latent_dim)
+
+    def call(self, inputs, training=None):
+        x = inputs
+
+        x = self.bottleneck(x)
+        x = self.flatten(x)
+        x = self.dense(x)
+        
+        return x
+
+
+class Decoder(tf.keras.Model):
+    
+    def __init__(self, bottleneck_size, latent_dim):
+        super(Decoder, self).__init__()
+
+        self.dense = tf.keras.layers.Dense(4 * 2 * bottleneck_size)
+        self.reshape = tf.keras.layers.Reshape((4, 2, bottleneck_size))
+        self.decoder = ConvNeXtReverse(depths=list(reversed([3, 3, 9, 3])), dims=list(reversed([bottleneck_size, 128, 64, 16])))
+
+    def call(self, inputs, training=None):
+        x = inputs
+
+        x = self.dense(x)
+        x = self.reshape(x)
+        x = self.decoder(x)
+
+        return x
+
+
+class AutoEncoder(tf.keras.Model):
+    
+    def __init__(self):
+        super(AutoEncoder, self).__init__()
+
+
+        bottleneck_size = 512
+        latent_dim = 128
+
+        self.encoder = ConvNeXt(depths=[3, 3, 9, 9, 3, 3], dims=[32, 64, 128, 256, 256, bottleneck_size], include_top=False)
+        self.bottleneck = Bottleneck(bottleneck_size, latent_dim)
+        self.decoder = Decoder(bottleneck_size, latent_dim)
+
+
+    def call(self, inputs, training=None):
+        x, _ = inputs
+
+        x = tf.transpose(x, [0,2,3,1])
+        x = self.encoder(x, training=training)
+        # x = self.pooling(x)
+        # x = self.reshape(x)
+        # x = self.flatten(x)
+
+        x = self.bottleneck(x, training=training)
+
+        x = self.decoder(x, training=training)
+        x = tf.image.resize(x, (512, 256), tf.image.ResizeMethod.LANCZOS5)
+        x = tf.transpose(x, [0,3,1,2])
+
+        # x = x * uv_mask[:,:, tf.newaxis]
+        x = x * uv_mask_body[tf.newaxis,tf.newaxis,:,:]
+
+        return x
+
+def draw_cicle(shape,diamiter):
+    '''
+    Input:
+    shape    : tuple (height, width)
+    diameter : scalar
+    
+    Output:
+    np.array of shape  that says True within a circle with diamiter =  around center 
+    '''
+    assert len(shape) == 2
+    TF = np.zeros(shape,dtype=bool)
+    center = np.array(TF.shape)/2.0
+
+    for iy in range(shape[0]):
+        for ix in range(shape[1]):
+            TF[iy,ix] = (iy- center[0])**2 + (ix - center[1])**2 < diamiter **2
+    return (TF)
+
+
+TFcircleIN   = draw_cicle(shape=(256,256), diamiter=15)
+TFcircleIN = np.ones((512,256,3), dtype=bool) * np.concatenate((TFcircleIN, TFcircleIN), axis=0)[:,:,np.newaxis]
+TFcircleOUT  = ~TFcircleIN
+
+def fft_loss_per_channel(pred, gt):
+    pred = tf.complex(pred, tf.zeros_like(pred))
+    gt = tf.complex(gt, tf.zeros_like(gt))
+
+    fft_pred_top = tf.stack([tf.signal.fftshift(tf.signal.fft2d(pred[:,i,:256,:])) for i in range(3)], axis=-1)
+    fft_gt_top = tf.stack([tf.signal.fftshift(tf.signal.fft2d(gt[:,i,:256,:])) for i in range(3)], axis=-1)
+    
+    fft_pred_bottom = tf.stack([tf.signal.fftshift(tf.signal.fft2d(pred[:,i,256:,:])) for i in range(3)], axis=-1)
+    fft_gt_bottom = tf.stack([tf.signal.fftshift(tf.signal.fft2d(gt[:,i,256:,:])) for i in range(3)], axis=-1)
+    
+    fft_pred = tf.concat((fft_pred_top, fft_pred_bottom), axis=1)
+    fft_gt = tf.concat((fft_gt_top, fft_gt_bottom), axis=1)
+
+    fft_pred_low = tf.where(TFcircleIN, fft_pred, tf.zeros_like(fft_pred))
+    fft_gt_low = tf.where(TFcircleIN, fft_gt, tf.zeros_like(fft_gt))
+    fft_pred_high = tf.where(TFcircleOUT, fft_pred, tf.zeros_like(fft_pred))
+    fft_gt_high = tf.where(TFcircleOUT, fft_gt, tf.zeros_like(fft_gt))
+
+    low_loss = tf.reduce_mean(tf.reduce_sum(tf.math.abs(fft_pred_low - fft_gt_low), axis=[1,2,3]))
+    high_loss = tf.reduce_mean(tf.reduce_sum(tf.math.abs(fft_pred_high - fft_gt_high), axis=[1,2,3]))
+    total = tf.reduce_mean(tf.reduce_sum(tf.math.abs(fft_pred - fft_gt), axis=[1,2,3]))
+
+    return total, low_loss, high_loss
+
+    # ex.imshow(tf.math.real(gt[0,0,:256,:])).show()
+    # img = tf.signal.fftshift(fft_gt[0][0])
+    # ex.imshow(tf.math.abs(tf.math.log(img))).show()
+    # low = tf.where(TFcircleIN, img, tf.zeros_like(img))
+    # ex.imshow(tf.math.real(tf.signal.ifft2d(tf.signal.ifftshift(low)))).show()
+    # high = tf.where(TFcircleOUT, img, tf.zeros_like(img))
+    # # ex.imshow(tf.math.abs(tf.math.log(high))).show()
+    # ex.imshow(tf.math.real(tf.signal.ifft2d(tf.signal.ifftshift(high)))).show()
+
+
+
+
+    print()
+
+
+@tf.function
+def train_step(batch):
+
+    with tf.GradientTape() as tape:
+
+        pred = model([batch['uv_body_posed'], batch['uv_mask']])
+        # pred = (pred + zero_body) * batch['mesh_mask'][:,:,tf.newaxis]
+
+        mesh_pred = uv2mesh(pred)
+        # mesh_pred = pred
+        mesh_orig = batch['mesh_body_posed']
+
+
+        e_loss = 1e-2 * tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']))
+        c_loss = 0 # 1e-2 * tf.reduce_mean(collision_penalty(mesh_pred, batch['mesh_body_posed'], batch['mesh_mask']))
+        # l1_loss = tf.reduce_mean(tf.reduce_sum((tf.math.abs(batch['mesh_unposed'] - pred)), axis=-1) * batch['mesh_mask'])
+        l1_loss = tf.reduce_mean(tf.reduce_sum((tf.math.abs(batch['uv_body_posed'] - pred)), axis=1) * batch['uv_mask'])
+        ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(tf.transpose(pred, [0,2,3,1]) + 1, tf.transpose(batch['uv_body_posed'], [0,2,3,1]) + 1, max_val=2))
+        rfft_loss = 0 #tf.reduce_mean(tf.reduce_sum(tf.math.abs(tf.signal.rfft3d(pred) - tf.signal.rfft3d(batch['uv_unposed'])), axis=[1,2,3]))
+        fft_loss = 0 # tf.reduce_mean(tf.reduce_sum(tf.math.abs(tf.signal.fft3d(tf.complex(pred, tf.zeros_like(pred))) - tf.signal.fft3d(tf.complex(batch['uv_body_posed'], tf.zeros_like(batch['uv_body_posed'])))), axis=[1,2,3]))
+        fft_loss_pc, fft_loss_pc_low, fft_loss_pc_high = 0,0,0 #fft_loss_per_channel(pred, batch['uv_unposed'])
+        # loss = fft_loss_pc # + e_loss + c_loss
+        loss = 1e3 * (1 - ms_ssim)
+
+    mesh_orig_2 = uv2mesh(batch['uv_body_posed'])
+    # mesh_skinned = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred)
+    metric = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
+    metric_2 = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig_2 - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
+    euclid_posed = 0 # tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_skinned - batch['mesh_posed']) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
+    gradients = tape.gradient(loss, model.trainable_variables)
+    # gradients = [tf.clip_by_norm(g, 1.0) for g in gradients]
+    # gradients = [tf.clip_by_value(g, -1.0, 1.0) for g in gradients]
+    optimizer.apply_gradients(zip(gradients, model.trainable_variables))
+    return { 'train/loss': loss, 'train/edge_loss': e_loss, 'train/collision_loss': c_loss, 'train/l1_loss': l1_loss, 'train/euclid_error': metric, 'train/ms_ssim': ms_ssim,
+             'train/fft_loss': fft_loss, 'train/euclid_error_reconst': metric_2, 'train/rfft_loss': rfft_loss, 'train/fft_oc_low': fft_loss_pc_low, 'train/fft_pc_high': fft_loss_pc_high,
+             'train/fft_loss_per_channel': fft_loss_pc, 'train/euclid_error_posed': euclid_posed }
+        
+
+
+@tf.function
+def valid_step(batch):
+    pred = model([batch['uv_body_posed'], batch['uv_mask']])
+
+    mesh_pred = uv2mesh(pred)
+    mesh_orig = batch['mesh_body_posed']
+    e_loss = 1e-2 * tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']))
+    c_loss = 0 #1e-2 * tf.reduce_mean(collision_penalty(mesh_pred, batch['mesh_body_posed'], batch['mesh_mask']))
+    l1_loss = tf.reduce_mean(tf.reduce_sum((tf.math.abs(batch['uv_body_posed'] - pred)), axis=1) * batch['uv_mask'])
+    ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(tf.transpose(pred, [0,2,3,1]) + 1, tf.transpose(batch['uv_body_posed'], [0,2,3,1]) + 1, max_val=2))
+    fft_loss = 0 #tf.reduce_mean(tf.reduce_sum(tf.math.abs(tf.signal.fft3d(tf.complex(pred, tf.zeros_like(pred))) - tf.signal.fft3d(tf.complex(batch['uv_body_posed'], tf.zeros_like(batch['uv_body_posed'])))), axis=[1,2,3]))
+    loss = 1e3 * (1 - ms_ssim)
+    mesh_orig_2 = uv2mesh(batch['uv_body_posed'])
+    metric = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
+    metric_2 = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig_2 - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
+  
+    return { 'valid/loss': loss, 'valid/edge_loss': e_loss, 'valid/collision_loss': c_loss, 'valid/l1_loss': l1_loss, 'valid/euclid_error': metric, 'valid/ms_ssim': ms_ssim,
+             'valid/fft_loss': fft_loss, 'valid/euclid_error_reconst': metric_2
+            }
+
+
+def plot_everything(batch: dict, type: str):
+
+    pred = model([batch['uv_body_posed'], batch['uv_mask']])
+    
+    mesh_pred = uv2mesh(pred)
+    mesh_orig = batch['mesh_body_posed']
+    
+    orig_img = tf.transpose(batch['uv_body_posed'], [0,2,3,1])
+    pred_img = tf.transpose(pred, [0,2,3,1])
+    img = tf.concat((orig_img, pred_img), axis=2)
+    imgs = [wandb.Image((img[i] + 1) / 2.0) for i in range(4)]
+
+    F = np.load("custom_utils/smpl_faces.npy")
+    mask = batch['mesh_mask'].numpy().astype(bool)
+    # garment_unposed, faces = maskFacesBatch(mesh_pred.numpy(), F, mask)
+    garment_unposed, faces = mesh_pred.numpy(), F
+    # gt = [mesh_orig[i].numpy()[mask[i], :] for i in range(pred.shape[0])]
+    gt = [mesh_orig[i].numpy()[:, :] for i in range(pred.shape[0])]
+    
+    body = batch['mesh_body_posed'].numpy()
+    # mesh_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred).numpy()
+    # garment_posed = [mesh_posed[i, mask[i], :] for i in range(body.shape[0])]
+    garment_posed = garment_unposed
+    collision = [np.concatenate((body[i], garment_posed[i]), axis=0) for i in range(len(garment_posed))]
+    # collision_faces = [np.concatenate((F, faces[i] + 14475), axis=0) for i in range(len(garment_posed))]
+    collision_faces = [np.concatenate((F, faces + 14475), axis=0) for i in range(len(garment_posed))]
+    collision_color = [np.concatenate((np.ones_like(body[i], dtype=np.int32)*100, np.ones_like(garment_posed[i], dtype=np.int32)*230), axis=0) for i in range(len(garment_posed))]
+
+    figs = { type + "/garment_" + str(i): plot_3d(gt[i], garment_unposed[i], faces, collision[i], collision_faces[i], collision_color[i]) for i in range(4) }
+    # figs = { type + "/garment_" + str(i): plot_3d(gt[i], garment_unposed[i], faces[i], collision[i], collision_faces[i], collision_color[i]) for i in range(4) }
+    
+    return { **figs, **{ type + '/images': imgs } }
+
+
+
+def train():
+    valid_iterator = iter(dataset_valid)
+
+    step = 0
+    for batch in dataset_train:
+        batch = static_dataset_train.expand_batch(batch_to_dict(batch, signature_dynamic))
+        train_loss = train_step(batch)
+
+        # data_valid = valid_iterator.get_next()
+        # data_valid = static_dataset_valid.expand_batch(batch_to_dict(data_valid, signature_dynamic))
+        valid_loss = {}#valid_step(data_valid)
+
+        if step % 1000 == 0:
+            train_figs = plot_everything(batch, 'train')
+            valid_figs = {}#plot_everything(data_valid, 'valid')
+            wandb.log({ **train_loss, **valid_loss, **train_figs, **valid_figs }, step=step)
+        else:
+            wandb.log({ **train_loss, **valid_loss }, step=step)
+            
+        if step % 2500 == 0:
+            manager.save(checkpoint_number=step)
+
+
+
+        print(step)
+
+        step += 1
+
+if __name__ == '__main__':
+    init_wandb()
+    allow_memory_growth()
+    parser = argparse.ArgumentParser(description='')
+    parser.add_argument('--name', type=str)
+    args = vars(parser.parse_args())
+    
+    smpl = SMPL()
+    uv_mask_body = np.load('custom_utils/mask_256x512.npy')
+    uv_mask_body = cv2.dilate(uv_mask_body.astype(np.float32), cv2.getStructuringElement(cv2.MORPH_RECT, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=2)
+
+    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_train = factory.get_dataset(folder='/home/workspace/local-data/jumpsuit_20/train', batch_size=32, signature=signature_dynamic)
+    static_dataset_train = StaticDataset(folder='/home/workspace/local-data/jumpsuit_20/train', signature=signature_static)
+    dataset_valid = factory.get_dataset(folder='/home/workspace/local-data/jumpsuit_20/valid', batch_size=32, signature=signature_dynamic)
+    static_dataset_valid = StaticDataset(folder='/home/workspace/local-data/jumpsuit_20/valid', signature=signature_static)
+
+    model = AutoEncoder()
+    # lr_schedule = tf.keras.optimizers.schedules.ExponentialDecay(0.0005, decay_steps=30000, decay_rate=0.1, staircase=False)
+    optimizer = tf.optimizers.Adam(learning_rate=0.0001)
+    folder = 'models/ae_logs/'
+
+    ckpt = tf.train.Checkpoint(model=model)
+    manager = tf.train.CheckpointManager(ckpt, 'checkpoints/' + name, max_to_keep=5)
+
+    train()

auto_encoder_curv.py

@@ -0,0 +1,379 @@
+import os
+import time
+import argparse
+import wandb
+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
+import plotly.express as ex
+
+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 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 convnext.models.convnext_tf import ConvNeXt, ConvNeXtReverse, Block, ConvNeXtReverseLabels
+from stylegan2_ae.generator import Generator
+from custom_utils.graph_utils import Reconstructor_V4_Decoder as GNNDecoder
+from custom_utils.graph_utils import Reconstructor_v2 as SimpleReconstructor
+from dataset import DatasetFactory, StaticDataset
+from custom_utils.custom_utils import batch_to_dict
+from custom_utils.smpl import SMPL
+
+name='ae_curvature_6'
+testing = False
+def init_wandb():
+    os.environ["CUDA_VISIBLE_DEVICES"] = "3"
+    if not testing:
+        run = wandb.init(project='vae', name=name, notes='convnext + convnext, 16*8*32, ce loss. flattened latent dim of 128')
+        artifact = wandb.Artifact('code', type='code')
+        artifact.add_file('auto_encoder_curv.py')
+        artifact.add_file('stylegan2_ae/layers/mapping_block.py')
+        artifact.add_file('convnext/models/convnext_tf.py')
+        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)
+    filtered_resolutions = resolutions[:index + 1]
+    filtered_featuremaps = featuremaps[:index + 1]
+    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 Bottleneck(tf.keras.Model):
+    
+    def __init__(self, dim, latent_dim):
+        super(Bottleneck, self).__init__()
+
+        self.bottleneck = tf.keras.Sequential(
+                [Block(dim=dim, drop_path=0,
+                       layer_scale_init_value=1e-6, prefix=f'bottleneck{j}') for j in range(5)]
+            )
+        self.flatten = tf.keras.layers.Flatten()
+        self.dense = tf.keras.layers.Dense(latent_dim)
+
+    def call(self, inputs, training=None):
+        x = inputs
+
+        x = self.bottleneck(x)
+        x = self.flatten(x)
+        x = self.dense(x)
+        
+        return x
+
+
+class Decoder(tf.keras.Model):
+    
+    def __init__(self, bottleneck_size, latent_dim):
+        super(Decoder, self).__init__()
+
+        self.dense = tf.keras.layers.Dense(16 * 8 * bottleneck_size)
+        self.reshape = tf.keras.layers.Reshape((16, 8, bottleneck_size))
+        self.decoder = ConvNeXtReverseLabels(depths=list(reversed([3, 3, 9, 3])), dims=list(reversed([bottleneck_size, 128, 64, 16])), nr_of_classes=7)
+
+    def call(self, inputs, training=None):
+        x = inputs
+
+        x = self.dense(x)
+        x = self.reshape(x)
+        x = self.decoder(x)
+
+        return x
+
+        
+class AutoEncoder(tf.keras.Model):
+    
+    def __init__(self):
+        super(AutoEncoder, self).__init__()
+
+
+        bottleneck_size = 32
+        latent_dim = 128
+        self.decoder = Decoder(bottleneck_size, latent_dim)
+
+        self.bottleneck = Bottleneck(bottleneck_size, latent_dim)
+
+        self.encoder = ConvNeXt(depths=[3, 3, 9, 3], dims=[32, 64, 128, bottleneck_size], include_top=False)
+
+
+
+    def call(self, inputs, training=None):
+        x, uv_mask = inputs
+
+        # x = tf.transpose(x, [0,2,3,1])
+        x = self.encoder(x, training=training)
+
+
+        x = self.bottleneck(x, training=training)
+
+        x = self.decoder(x, training=training)
+        # x = tf.transpose(x, [0,3,1,2])
+
+        # x = x * uv_mask[:,:, tf.newaxis]
+        # x = x * uv_mask[:,tf.newaxis,:,:]
+        x = x * uv_mask[:,:,:,tf.newaxis]
+
+        return x
+
+
+
+
+g_mean_iou = tf.keras.metrics.MeanIoU(num_classes=7)
+def get_mean_iou(y_true, y_pred):
+    g_mean_iou.reset_state()
+    g_mean_iou.update_state(y_true, y_pred)
+    return g_mean_iou.result()
+
+# https://github.com/NifTK/NiftyNet/blob/dev/niftynet/layer/loss_segmentation.py
+def dice_dense(prediction, ground_truth, weight_map=None):
+    """
+    Computing mean-class Dice similarity.
+    :param prediction: last dimension should have ``num_classes``
+    :param ground_truth: segmentation ground truth (encoded as a binary matrix)
+        last dimension should be ``num_classes``
+    :param weight_map:
+    :return: ``1.0 - mean(Dice similarity per class)``
+    """
+
+    if weight_map is not None:
+        raise NotImplementedError
+    prediction = tf.cast(prediction, dtype=tf.float32)
+    ground_truth = tf.cast(ground_truth, dtype=tf.float32)
+    ground_truth = tf.reshape(ground_truth, prediction.shape)
+    # computing Dice over the spatial dimensions
+    reduce_axes = list(range(len(prediction.shape) - 1))
+    dice_numerator = 2.0 * tf.reduce_sum(
+        prediction * ground_truth, axis=reduce_axes)
+    dice_denominator = \
+        tf.reduce_sum(tf.square(prediction), axis=reduce_axes) + \
+        tf.reduce_sum(tf.square(ground_truth), axis=reduce_axes)
+
+    epsilon = 0.00001
+
+    dice_score = (dice_numerator + epsilon) / (dice_denominator + epsilon)
+    return tf.reduce_mean(dice_score)
+
+
+@tf.function
+def train_step(batch):
+
+    with tf.GradientTape() as tape:
+        uv_curvature = tf.one_hot(batch['uv_curvature'], 7)
+
+        pred = model([uv_curvature, batch['uv_mask']])
+
+
+        ce_uv_loss = scce(batch['uv_curvature'], pred)
+        ms_ssim = tf.reduce_mean(tf.image.ssim(pred, uv_curvature, max_val=1))
+        # ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(tf.math.argmax(pred, axis=-1, output_type=tf.dtypes.int32)[:,:,:,tf.newaxis], batch['uv_curvature'][:,:,:,tf.newaxis], max_val=7*2))
+
+        # loss =  1e3 * (1 - ms_ssim)
+        loss = ce_uv_loss
+
+    # mesh_pred = uv2mesh(tf.tile(tf.math.argmax(pred, axis=-1, output_type=tf.dtypes.int32)[:,:,:,tf.newaxis], [1,1,1,3]))[:,:,0]
+    # mesh_orig = tf.one_hot(batch['mesh_curvature'] + tf.cast(batch['mesh_mask'], tf.int32), 7)
+    # ce_mesh_loss = scce(mesh_orig, mesh_pred)
+
+    mean_iou = get_mean_iou(batch['uv_curvature'], tf.math.argmax(pred, axis=-1))
+    dice = dice_dense(pred, uv_curvature)
+
+    gradients = tape.gradient(loss, model.trainable_variables)
+    optimizer.apply_gradients(zip(gradients, model.trainable_variables))
+    return { 'train/loss': loss, 'train/ce_uv_loss': ce_uv_loss, 'train/ssim': ms_ssim, 'train/mean_iou': mean_iou, 'train/dice': dice  #'ce_mesh_loss': ce_mesh_loss
+             }
+        
+
+
+@tf.function
+def valid_step(batch):
+    uv_curvature = tf.one_hot(batch['uv_curvature'], 7)
+
+    pred = model([uv_curvature, batch['uv_mask']])
+
+    ce_uv_loss = scce(batch['uv_curvature'], pred)
+    ms_ssim = tf.reduce_mean(tf.image.ssim(pred, uv_curvature, max_val=1))
+    loss =  ce_uv_loss
+
+    mean_iou = get_mean_iou(batch['uv_curvature'], tf.math.argmax(pred, axis=-1))
+    dice = dice_dense(pred, uv_curvature)
+
+    return { 'valid/loss': loss, 'valid/ce_uv_loss': ce_uv_loss, 'valid/ssim': ms_ssim, 'valid/mean_iou': mean_iou, 'valid/dice': dice #'ce_mesh_loss': ce_mesh_loss
+             }
+
+
+
+def train():
+    step = 0
+    for batch in dataset_train:
+        batch = static_dataset_train.expand_batch(batch_to_dict(batch, signature_dynamic))
+
+        train_loss = train_step(batch)
+
+        data_valid = valid_iterator.get_next()
+        data_valid = static_dataset_valid.expand_batch(batch_to_dict(data_valid, signature_dynamic))
+        valid_loss = valid_step(data_valid)
+
+        to_log = { **train_loss, **valid_loss}
+        if step % 1000 == 0:
+            uv_curvature = tf.one_hot(data_valid['uv_curvature'], 7)
+            pred = model([uv_curvature, data_valid['uv_mask']])
+            
+            orig_img = data_valid['uv_curvature']
+            pred_img = tf.math.argmax(pred, axis=-1, output_type=tf.dtypes.int32)
+            img = tf.concat((orig_img, pred_img), axis=2).numpy()
+            imgs = { 'valid/images': [wandb.Image(img[i,:,:,np.newaxis] / 6.0) for i in range(4)] }
+            to_log = { **to_log, **imgs}
+
+        wandb.log(to_log, step=step)  
+        print(step)
+
+        if step % 2500 == 0:
+            manager.save(checkpoint_number=step)
+
+        step += 1
+
+
+def test():
+    ckpt.restore(manager.latest_checkpoint)
+    for batch in dataset_train:
+        batch = static_dataset_train.expand_batch(batch_to_dict(batch, signature_dynamic))
+        uv_curvature = tf.one_hot(batch['uv_curvature'], 7)
+
+        pred = model([uv_curvature, batch['uv_mask']])
+        orig_img = batch['uv_curvature']
+        pred_img = tf.math.argmax(pred, axis=-1, output_type=tf.dtypes.int32)
+        img = tf.concat((orig_img, pred_img), axis=2).numpy()
+
+        ex.imshow(img[0] / 6.0).show()
+        print()
+
+
+
+if __name__ == '__main__':
+    init_wandb()
+    allow_memory_growth()
+    parser = argparse.ArgumentParser(description='')
+    parser.add_argument('--name', type=str)
+    args = vars(parser.parse_args())
+    
+    smpl = SMPL()
+    scce = tf.keras.losses.SparseCategoricalCrossentropy()
+
+    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,)
+                          'uv_curvature': (512, 256),
+                          'mesh_curvature':  (14475)
+                        }
+    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_train = factory.get_dataset(folder='/home/workspace/local-data/jumpsuit_20/train',
+                                  batch_size=8,
+                                  signature=signature_dynamic)
+    static_dataset_train = StaticDataset(folder='/home/workspace/local-data/jumpsuit_20/train',
+                                   signature=signature_static)
+    dataset_valid = factory.get_dataset(folder='/home/workspace/local-data/jumpsuit_20/valid',
+                                  batch_size=8,
+                                  signature=signature_dynamic)
+    static_dataset_valid = StaticDataset(folder='/home/workspace/local-data/jumpsuit_20/valid',
+                                   signature=signature_static)
+
+    valid_iterator = iter(dataset_valid)
+
+    model = AutoEncoder()
+    # optimizer = tf.optimizers.Adam(learning_rate=0.002, beta_1=0.0, beta_2=0.99, epsilon=1e-8)
+    lr_schedule = tf.keras.optimizers.schedules.ExponentialDecay(0.001, decay_steps=10000, decay_rate=0.1, staircase=False)
+    optimizer = tf.optimizers.Adam(learning_rate=lr_schedule)
+    folder = 'models/ae_logs/'
+
+    ckpt = tf.train.Checkpoint(model=model)
+    manager = tf.train.CheckpointManager(ckpt, 'checkpoints/' + name, max_to_keep=5)
+
+    if testing:
+        test()
+    else:
+        train()

auto_encoder_static.py

@@ -0,0 +1,450 @@
+import os
+import time
+import argparse
+import wandb
+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
+import plotly.express as ex
+
+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 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 convnext.models.convnext_tf import ConvNeXt, ConvNeXtReverse, Block
+from stylegan2_ae.generator import Generator
+from custom_utils.graph_utils import Reconstructor_V4_Decoder as GNNDecoder
+from custom_utils.graph_utils import Reconstructor_v2 as SimpleReconstructor
+from dataset import DatasetFactory, StaticDataset
+from custom_utils.custom_utils import batch_to_dict
+from custom_utils.smpl import SMPL
+
+name='ae_static_3'
+def init_wandb():
+    os.environ["CUDA_VISIBLE_DEVICES"] = "2"
+    run = wandb.init(project='vae', name=name,
+            notes='convnext + convnext, add dense with tanh at the end of reverse convnext, larger features. add lr scheduler. fft3d instead of rfft3d. \
+    fft2d per channel and separate for top and bottom. fft_loss, all data. added convnext block as bottleneck, same as before, added extra metric. \
+    . lowered learning rate because it goes to nan. shallow third stage. slower lr decay. with validation. higher lr rate.\
+    convnext first stage down/upscale by factor fo 2 instead of 4, and 5 stages instead of 4 so bottlneck same size\
+    flattened latent dim to 128')
+    artifact = wandb.Artifact('code', type='code')
+    artifact.add_file('auto_encoder_static.py')
+    artifact.add_file('losses.py')
+    artifact.add_file('dataset.py')
+    artifact.add_file('convnext/models/convnext_tf.py')
+    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)
+    filtered_resolutions = resolutions[:index + 1]
+    filtered_featuremaps = featuremaps[:index + 1]
+    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
+
+
+def plot_3d(mesh_gt, mesh_pred_unposed, faces, collision, collision_faces, collision_color):
+    mesh_pred_unposed[:,1] += 1
+    collision[:,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_unposed[:,0],
+            y=mesh_pred_unposed[:,1],
+            z=mesh_pred_unposed[:,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=collision[:,0],
+            y=collision[:,1],
+            z=collision[:,2],
+            # i, j and k give the vertices of triangles
+            i = collision_faces[:,0],
+            j = collision_faces[:,1],
+            k = collision_faces[:,2],
+            vertexcolor=collision_color
+        )
+    ])
+
+    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 Bottleneck(tf.keras.Model):
+    
+    def __init__(self, dim, latent_dim):
+        super(Bottleneck, self).__init__()
+
+        self.bottleneck = tf.keras.Sequential(
+                [Block(dim=dim, drop_path=0,
+                       layer_scale_init_value=1e-6, prefix=f'bottleneck{j}') for j in range(5)]
+            )
+        self.flatten = tf.keras.layers.Flatten()
+        self.dense = tf.keras.layers.Dense(latent_dim)
+
+    def call(self, inputs, training=None):
+        x = inputs
+
+        x = self.bottleneck(x)
+        x = self.flatten(x)
+        x = self.dense(x)
+        
+        return x
+
+
+class Decoder(tf.keras.Model):
+    
+    def __init__(self, bottleneck_size, latent_dim):
+        super(Decoder, self).__init__()
+
+        self.dense = tf.keras.layers.Dense(8 * 4 * bottleneck_size)
+        self.reshape = tf.keras.layers.Reshape((8, 4, bottleneck_size))
+        self.decoder = ConvNeXtReverse(depths=list(reversed([3, 3, 9, 3, 3])), dims=list(reversed([bottleneck_size, 128, 128, 64, 16])))
+
+    def call(self, inputs, training=None):
+        x = inputs
+
+        x = self.dense(x)
+        x = self.reshape(x)
+        x = self.decoder(x)
+
+        return x
+
+
+class AutoEncoder(tf.keras.Model):
+    
+    def __init__(self):
+        super(AutoEncoder, self).__init__()
+
+
+        bottleneck_size = 32
+        latent_dim = 128
+
+        self.encoder = ConvNeXt(depths=[3, 3, 9, 3, 3], dims=[32, 64, 128, 128, bottleneck_size], include_top=False)
+        self.bottleneck = Bottleneck(bottleneck_size, latent_dim)
+        self.decoder = Decoder(bottleneck_size, latent_dim)
+
+    def call(self, inputs, training=None):
+        x, uv_mask = inputs
+
+        x = tf.transpose(x, [0,2,3,1])
+        x = self.encoder(x, training=training)
+        # x = self.pooling(x)
+        # x = self.reshape(x)
+        # x = self.flatten(x)
+
+        x = self.bottleneck(x, training=training)
+
+        x = self.decoder(x, training=training)
+        x = tf.transpose(x, [0,3,1,2])
+
+        # x = x * uv_mask[:,:, tf.newaxis]
+        x = x * uv_mask[:,tf.newaxis,:,:]
+
+        return x
+
+def draw_cicle(shape,diamiter):
+    '''
+    Input:
+    shape    : tuple (height, width)
+    diameter : scalar
+    
+    Output:
+    np.array of shape  that says True within a circle with diamiter =  around center 
+    '''
+    assert len(shape) == 2
+    TF = np.zeros(shape,dtype=bool)
+    center = np.array(TF.shape)/2.0
+
+    for iy in range(shape[0]):
+        for ix in range(shape[1]):
+            TF[iy,ix] = (iy- center[0])**2 + (ix - center[1])**2 < diamiter **2
+    return (TF)
+
+
+TFcircleIN   = draw_cicle(shape=(256,256), diamiter=15)
+TFcircleIN = np.ones((512,256,3), dtype=bool) * np.concatenate((TFcircleIN, TFcircleIN), axis=0)[:,:,np.newaxis]
+TFcircleOUT  = ~TFcircleIN
+
+def fft_loss_per_channel(pred, gt):
+    pred = tf.complex(pred, tf.zeros_like(pred))
+    gt = tf.complex(gt, tf.zeros_like(gt))
+
+    fft_pred_top = tf.stack([tf.signal.fftshift(tf.signal.fft2d(pred[:,i,:256,:])) for i in range(3)], axis=-1)
+    fft_gt_top = tf.stack([tf.signal.fftshift(tf.signal.fft2d(gt[:,i,:256,:])) for i in range(3)], axis=-1)
+    
+    fft_pred_bottom = tf.stack([tf.signal.fftshift(tf.signal.fft2d(pred[:,i,256:,:])) for i in range(3)], axis=-1)
+    fft_gt_bottom = tf.stack([tf.signal.fftshift(tf.signal.fft2d(gt[:,i,256:,:])) for i in range(3)], axis=-1)
+    
+    fft_pred = tf.concat((fft_pred_top, fft_pred_bottom), axis=1)
+    fft_gt = tf.concat((fft_gt_top, fft_gt_bottom), axis=1)
+
+    fft_pred_low = tf.where(TFcircleIN, fft_pred, tf.zeros_like(fft_pred))
+    fft_gt_low = tf.where(TFcircleIN, fft_gt, tf.zeros_like(fft_gt))
+    fft_pred_high = tf.where(TFcircleOUT, fft_pred, tf.zeros_like(fft_pred))
+    fft_gt_high = tf.where(TFcircleOUT, fft_gt, tf.zeros_like(fft_gt))
+
+    low_loss = tf.reduce_mean(tf.reduce_sum(tf.math.abs(fft_pred_low - fft_gt_low), axis=[1,2,3]))
+    high_loss = tf.reduce_mean(tf.reduce_sum(tf.math.abs(fft_pred_high - fft_gt_high), axis=[1,2,3]))
+    total = tf.reduce_mean(tf.reduce_sum(tf.math.abs(fft_pred - fft_gt), axis=[1,2,3]))
+
+    return total, low_loss, high_loss
+
+    # ex.imshow(tf.math.real(gt[0,0,:256,:])).show()
+    # img = tf.signal.fftshift(fft_gt[0][0])
+    # ex.imshow(tf.math.abs(tf.math.log(img))).show()
+    # low = tf.where(TFcircleIN, img, tf.zeros_like(img))
+    # ex.imshow(tf.math.real(tf.signal.ifft2d(tf.signal.ifftshift(low)))).show()
+    # high = tf.where(TFcircleOUT, img, tf.zeros_like(img))
+    # # ex.imshow(tf.math.abs(tf.math.log(high))).show()
+    # ex.imshow(tf.math.real(tf.signal.ifft2d(tf.signal.ifftshift(high)))).show()
+
+
+
+
+    print()
+
+
+@tf.function
+def train_step(batch):
+
+    with tf.GradientTape() as tape:
+        pred = model([batch['uv_static'], batch['uv_mask']])
+
+        mesh_pred = uv2mesh(pred)
+        mesh_orig = batch['mesh_static']
+
+        e_loss = 1e-2 * tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']))
+        c_loss = 0 #1e-2 * tf.reduce_mean(collision_penalty(mesh_pred, batch['mesh_body_posed'], batch['mesh_mask']))
+        # l1_loss = tf.reduce_mean(tf.reduce_sum((tf.math.abs(batch['mesh_unposed'] - pred)), axis=-1) * batch['mesh_mask'])
+        l1_loss = tf.reduce_mean(tf.reduce_sum((tf.math.abs(batch['uv_static'] - pred)), axis=1) * batch['uv_mask'])
+        ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(tf.transpose(pred, [0,2,3,1]) + 1, tf.transpose(batch['uv_static'], [0,2,3,1]) + 1, max_val=2))
+        fft_loss = tf.reduce_mean(tf.reduce_sum(tf.math.abs(tf.signal.fft3d(tf.complex(pred, tf.zeros_like(pred))) - tf.signal.fft3d(tf.complex(batch['uv_static'], tf.zeros_like(batch['uv_static'])))), axis=[1,2,3]))
+        fft_loss_pc, fft_loss_pc_low, fft_loss_pc_high = 0,0,0 #fft_loss_per_channel(pred, batch['uv_unposed'])
+        loss = 1e3 * (1 - ms_ssim)
+
+    mesh_orig_2 = uv2mesh(batch['uv_static'])
+    # mesh_skinned = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred)
+    metric = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
+    metric_2 = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig_2 - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
+    euclid_posed = 0 # tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_skinned - batch['mesh_posed']) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
+    gradients = tape.gradient(loss, model.trainable_variables)
+    # gradients = [tf.clip_by_norm(g, 1.0) for g in gradients]
+    # gradients = [tf.clip_by_value(g, -1.0, 1.0) for g in gradients]
+    optimizer.apply_gradients(zip(gradients, model.trainable_variables))
+    return { 'train/loss': loss, 'train/edge_loss': e_loss, 'train/collision_loss': c_loss, 'train/l1_loss': l1_loss, 'train/euclid_error': metric, 'train/ms_ssim': ms_ssim,
+             'train/fft_loss': fft_loss, 'train/euclid_error_reconst': metric_2, 'train/fft_oc_low': fft_loss_pc_low, 'train/fft_pc_high': fft_loss_pc_high,
+             'train/fft_loss_per_channel': fft_loss_pc, 'train/euclid_error_posed': euclid_posed }
+        
+
+
+@tf.function
+def valid_step(batch):
+    pred = model([batch['uv_static'], batch['uv_mask']])
+
+    mesh_pred = uv2mesh(pred)
+    mesh_orig = batch['mesh_static']
+    e_loss = 1e-2 * tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']))
+    c_loss = 0
+    l1_loss = tf.reduce_mean(tf.reduce_sum((tf.math.abs(batch['uv_static'] - pred)), axis=1) * batch['uv_mask'])
+    ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(tf.transpose(pred, [0,2,3,1]) + 1, tf.transpose(batch['uv_static'], [0,2,3,1]) + 1, max_val=2))
+    fft_loss = 0 #tf.reduce_mean(tf.reduce_sum(tf.math.abs(tf.signal.fft3d(tf.complex(pred, tf.zeros_like(pred))) - tf.signal.fft3d(tf.complex(batch['uv_static'], tf.zeros_like(batch['uv_static'])))), axis=[1,2,3]))
+    loss = 1e3 * (1 - ms_ssim)
+    mesh_orig_2 = uv2mesh(batch['uv_static'])
+    metric = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
+    metric_2 = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig_2 - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
+  
+    return { 'valid/loss': loss, 'valid/edge_loss': e_loss, 'valid/collision_loss': c_loss, 'valid/l1_loss': l1_loss, 'valid/euclid_error': metric, 'valid/ms_ssim': ms_ssim,
+             'valid/fft_loss': fft_loss, 'valid/euclid_error_reconst': metric_2
+            }
+
+
+def plot_everything(batch: dict, type: str):
+
+    pred = model([batch['uv_static'], batch['uv_mask']])
+    
+    mesh_pred = uv2mesh(pred)
+    mesh_orig = batch['mesh_static']
+    
+    orig_img = tf.transpose(batch['uv_static'], [0,2,3,1])
+    pred_img = tf.transpose(pred, [0,2,3,1])
+    img = tf.concat((orig_img, pred_img), axis=2)
+    imgs = [wandb.Image((img[i] + 1) / 2.0) for i in range(4)]
+
+    F = np.load("custom_utils/smpl_faces.npy")
+    mask = batch['mesh_mask'].numpy().astype(bool)
+    garment_unposed, faces = maskFacesBatch(mesh_pred.numpy(), F, mask)
+    gt = [mesh_orig[i].numpy()[mask[i], :] for i in range(pred.shape[0])]
+    
+    body = gt
+    garment_posed = mesh_pred
+    collision = [np.concatenate((body[i], garment_posed[i]), axis=0) for i in range(len(garment_posed))]
+    collision_faces = [np.concatenate((faces[i], faces[i] + np.sum(mask[i])), axis=0) for i in range(len(garment_posed))]
+    collision_color = [np.concatenate((np.ones_like(body[i], dtype=np.int32)*100, np.ones_like(garment_posed[i], dtype=np.int32)*230), axis=0) for i in range(len(garment_posed))]
+
+    figs = { type + "/garment_" + str(i): plot_3d(gt[i], garment_unposed[i], faces[i], collision[i], collision_faces[i], collision_color[i]) for i in range(4) }
+    
+    return { **figs, **{ type + '/images': imgs } }
+
+
+def train():
+    step = 0
+    for batch in dataset_train:
+        batch = static_dataset_train.expand_batch(batch_to_dict(batch, signature_dynamic))
+        train_loss = train_step(batch)
+
+        data_valid = valid_iterator.get_next()
+        data_valid = static_dataset_valid.expand_batch(batch_to_dict(data_valid, signature_dynamic))
+        valid_loss = valid_step(data_valid)
+
+        if step % 1000 == 0:
+            train_figs = plot_everything(batch, 'train')
+            valid_figs = plot_everything(data_valid, 'valid')
+            wandb.log({ **train_loss, **valid_loss, **train_figs, **valid_figs }, step=step)
+        else:
+            wandb.log({ **train_loss, **valid_loss }, step=step)
+            
+
+        print(step)
+
+        step += 1
+
+if __name__ == '__main__':
+    init_wandb()
+    allow_memory_growth()
+    parser = argparse.ArgumentParser(description='')
+    parser.add_argument('--name', type=str)
+    args = vars(parser.parse_args())
+    
+    smpl = SMPL()
+
+
+    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_train = factory.get_dataset(folder='/home/workspace/local-data/jumpsuit_20/train', batch_size=8, signature=signature_dynamic)
+    static_dataset_train = StaticDataset(folder='/home/workspace/local-data/jumpsuit_20/train', signature=signature_static)
+    dataset_valid = factory.get_dataset(folder='/home/workspace/local-data/jumpsuit_20/valid', batch_size=8, signature=signature_dynamic)
+    static_dataset_valid = StaticDataset(folder='/home/workspace/local-data/jumpsuit_20/valid', signature=signature_static)
+
+
+    model = AutoEncoder()
+    lr_schedule = tf.keras.optimizers.schedules.ExponentialDecay(0.0005, decay_steps=30000, decay_rate=0.1, staircase=False)
+    optimizer = tf.optimizers.Adam(learning_rate=lr_schedule)
+    folder = 'models/ae_logs/'
+    valid_iterator = iter(dataset_valid)
+
+    ckpt = tf.train.Checkpoint(model=model)
+    manager = tf.train.CheckpointManager(ckpt, 'checkpoints/' + name, max_to_keep=5)
+
+    train()

convnext/models/convnext_tf.py

@@ -1,9 +1,9 @@
+from unittest import skip
 import numpy as np
 import tensorflow as tf
-from tensorflow.keras.layers import (Conv2D, Dense, DepthwiseConv2D,
+from tensorflow.keras.layers import (Conv2D, Conv2DTranspose, Dense, DepthwiseConv2D,
                                      GlobalAveragePooling2D, Layer,
-                                     LayerNormalization)
+                                     LayerNormalization, Flatten, Reshape, UpSampling2D)
-
 
 def drop_path(inputs, drop_prob, is_training):
     # borrowed from https://github.com/rishigami/Swin-Transformer-TF/blob/main/swintransformer/model.py
@@ -109,6 +109,299 @@ class ConvNeXt(tf.keras.Model):
             LayerNormalization(epsilon=1e-6)]
         )
         self.downsample_layers.append(stem)
+        for i in range(len(depths) - 1):
+            downsample_layer = tf.keras.Sequential([
+                LayerNormalization(epsilon=1e-6),
+                Conv2D(dims[i+1], kernel_size=2, strides=2, padding='same')]
+            )
+            self.downsample_layers.append(downsample_layer)
+
+        self.stages = []  # 4 feature resolution stages, each consisting of multiple residual blocks
+        dp_rates = [x for x in np.linspace(0, drop_path_rate, sum(depths))]
+        cur = 0
+        for i in range(len(depths)):
+            stage = tf.keras.Sequential(
+                [Block(dim=dims[i], drop_path=dp_rates[cur + j],
+                       layer_scale_init_value=layer_scale_init_value, prefix=f'block{i}') for j in range(depths[i])]
+            )
+            self.stages.append(stage)
+            cur += depths[i]
+
+        if self.include_top:
+            self.avg = GlobalAveragePooling2D()
+            self.norm = LayerNormalization(epsilon=1e-6)  # final norm layer
+            self.head = Dense(num_classes)
+        else:
+            self.avg = None
+            self.norm = None
+            self.head = None
+
+    def forward_features(self, x):
+        for i in range(len(self.downsample_layers)):
+            x = self.downsample_layers[i](x)
+            x = self.stages[i](x)
+        return x
+
+    def call(self, x):
+        x = self.forward_features(x)
+        if self.include_top:
+            x = self.avg(x)
+            x = self.norm(x)
+            x = self.head(x)
+        return x
+
+
+class ConvNeXtWithSkip(tf.keras.Model):
+    r""" ConvNeXt
+        A Tensorflow keras impl of : `A ConvNet for the 2020s`  -
+          https://arxiv.org/pdf/2201.03545.pdf
+
+    Args:
+        num_classes (int): Number of classes for classification head. Default: 1000
+        depths (tuple(int)): Number of blocks at each stage. Default: [3, 3, 9, 3]
+        dims (int): Feature dimension at each stage. Default: [96, 192, 384, 768]
+        include_top (bool): whether to add head or just use it as feature extractor. Default: True
+        drop_path_rate (float): Stochastic depth rate. Default: 0.
+        layer_scale_init_value (float): Init value for Layer Scale. Default: 1e-6.
+        head_init_scale (float): Init scaling value for classifier weights and biases. Default: 1.
+    """
+
+    def __init__(self, 
+                 depths=[3, 3, 9, 3], dims=[96, 192, 384, 768], 
+                 drop_path_rate=0., layer_scale_init_value=1e-6, 
+                 ):
+        super().__init__()
+        self.downsample_layers = []  # stem and 3 intermediate downsampling conv layers
+        stem = tf.keras.Sequential([
+            Conv2D(dims[0], kernel_size=4, strides=4, padding='same'),
+            LayerNormalization(epsilon=1e-6)]
+        )
+        self.downsample_layers.append(stem)
+        for i in range(len(depths) - 1):
+            downsample_layer = tf.keras.Sequential([
+                LayerNormalization(epsilon=1e-6),
+                Conv2D(dims[i+1], kernel_size=2, strides=2, padding='same')]
+            )
+            self.downsample_layers.append(downsample_layer)
+
+        self.stages = []  # 4 feature resolution stages, each consisting of multiple residual blocks
+        dp_rates = [x for x in np.linspace(0, drop_path_rate, sum(depths))]
+        cur = 0
+        for i in range(len(depths)):
+            stage = tf.keras.Sequential(
+                [Block(dim=dims[i], drop_path=dp_rates[cur + j],
+                       layer_scale_init_value=layer_scale_init_value, prefix=f'block{i}') for j in range(depths[i])]
+            )
+            self.stages.append(stage)
+            cur += depths[i]
+
+
+    def call(self, x):
+        skip_features = []
+        for i in range(len(self.downsample_layers)):
+            x = self.downsample_layers[i](x)
+            x = self.stages[i](x)
+            skip_features.append(x)
+        return x, skip_features
+
+
+class ConvNeXtReverse(tf.keras.Model):
+    r""" ConvNeXt
+        A Tensorflow keras impl of : `A ConvNet for the 2020s`  -
+          https://arxiv.org/pdf/2201.03545.pdf
+
+    Args:
+        num_classes (int): Number of classes for classification head. Default: 1000
+        depths (tuple(int)): Number of blocks at each stage. Default: [3, 3, 9, 3]
+        dims (int): Feature dimension at each stage. Default: [96, 192, 384, 768]
+        include_top (bool): whether to add head or just use it as feature extractor. Default: True
+        drop_path_rate (float): Stochastic depth rate. Default: 0.
+        layer_scale_init_value (float): Init value for Layer Scale. Default: 1e-6.
+        head_init_scale (float): Init scaling value for classifier weights and biases. Default: 1.
+    """
+
+    def __init__(self, 
+                 depths=[3, 3, 9, 3], dims=[96, 192, 384, 768],
+                 drop_path_rate=0., layer_scale_init_value=1e-6, head_init_scale=1.,
+                 ):
+        super().__init__()
+        self.upsample_layers = []  # stem and 3 intermediate downsampling conv layers
+        stem = tf.keras.Sequential([
+            Conv2DTranspose(dims[0], kernel_size=4, strides=4, padding='same'),
+            # UpSampling2D((4, 4), interpolation='nearest'),
+            # Conv2D(dims[0], kernel_size=4, strides=1, padding='same'),
+            # Dense(dims[0]),
+            LayerNormalization(epsilon=1e-6)]
+        )
+        self.upsample_layers.append(stem)
+        for i in range(len(depths) - 1):
+            downsample_layer = tf.keras.Sequential([
+                LayerNormalization(epsilon=1e-6),
+                Conv2DTranspose(dims[i+1], kernel_size=2, strides=2, padding='same')
+                # UpSampling2D((2,2), interpolation='nearest'),
+                # Conv2D(dims[i+1], kernel_size=2, strides=1, padding='same'),
+                # Dense(dims[i+1]),
+            ])
+            self.upsample_layers.append(downsample_layer)
+
+        self.stages = []  # 4 feature resolution stages, each consisting of multiple residual blocks
+        dp_rates = [x for x in np.linspace(0, drop_path_rate, sum(depths))]
+        cur = 0
+        for i in range(len(depths)):
+            stage = tf.keras.Sequential(
+                [Block(dim=dims[i], drop_path=dp_rates[cur + j],
+                       layer_scale_init_value=layer_scale_init_value, prefix=f'block{i}') for j in range(depths[i])]
+            )
+            self.stages.append(stage)
+            cur += depths[i]
+
+        self.last = Dense(3, activation='tanh')
+
+    def call(self, x):
+        for i in reversed(range(len(self.upsample_layers))):
+            x = self.upsample_layers[i](x)
+            x = self.stages[i](x)
+        x = self.last(x)
+        return x
+
+
+class ConvNeXtReverseWithSkip(tf.keras.Model):
+    r""" ConvNeXt
+        A Tensorflow keras impl of : `A ConvNet for the 2020s`  -
+          https://arxiv.org/pdf/2201.03545.pdf
+
+    Args:
+        num_classes (int): Number of classes for classification head. Default: 1000
+        depths (tuple(int)): Number of blocks at each stage. Default: [3, 3, 9, 3]
+        dims (int): Feature dimension at each stage. Default: [96, 192, 384, 768]
+        include_top (bool): whether to add head or just use it as feature extractor. Default: True
+        drop_path_rate (float): Stochastic depth rate. Default: 0.
+        layer_scale_init_value (float): Init value for Layer Scale. Default: 1e-6.
+        head_init_scale (float): Init scaling value for classifier weights and biases. Default: 1.
+    """
+
+    def __init__(self, 
+                 depths=[3, 3, 9, 3], dims=[96, 192, 384, 768],
+                 drop_path_rate=0., layer_scale_init_value=1e-6, head_init_scale=1.,
+                 ):
+        super().__init__()
+        assert len(depths) == len(dims)
+        self.upsample_layers = []  # stem and 3 intermediate downsampling conv layers
+        stem = tf.keras.Sequential([
+            Conv2DTranspose(dims[0], kernel_size=4, strides=4, padding='same'),
+            LayerNormalization(epsilon=1e-6)]
+        )
+        self.upsample_layers.append(stem)
+        for i in range(len(depths) - 1):
+            downsample_layer = tf.keras.Sequential([
+                LayerNormalization(epsilon=1e-6),
+                Conv2DTranspose(dims[i+1], kernel_size=2, strides=2, padding='same')]
+            )
+            self.upsample_layers.append(downsample_layer)
+
+        self.stages = []  # 4 feature resolution stages, each consisting of multiple residual blocks
+        dp_rates = [x for x in np.linspace(0, drop_path_rate, sum(depths))]
+        cur = 0
+        for i in range(len(depths)):
+            stage = tf.keras.Sequential(
+                [Block(dim=dims[i], drop_path=dp_rates[cur + j],
+                       layer_scale_init_value=layer_scale_init_value, prefix=f'block{i}') for j in range(depths[i])]
+            )
+            self.stages.append(stage)
+            cur += depths[i]
+
+        self.last = Dense(3, activation='tanh')
+
+    def call(self, x):
+        x, skip_features = x
+        for i in reversed(range(len(self.upsample_layers))):
+            x = tf.concat((x, skip_features[i]), axis=-1)
+            x = self.upsample_layers[i](x)
+            x = self.stages[i](x)
+        x = self.last(x)
+        return x
+
+
+class ConvNeXtReverseLabels(tf.keras.Model):
+    r""" ConvNeXt
+        A Tensorflow keras impl of : `A ConvNet for the 2020s`  -
+          https://arxiv.org/pdf/2201.03545.pdf
+
+    Args:
+        num_classes (int): Number of classes for classification head. Default: 1000
+        depths (tuple(int)): Number of blocks at each stage. Default: [3, 3, 9, 3]
+        dims (int): Feature dimension at each stage. Default: [96, 192, 384, 768]
+        include_top (bool): whether to add head or just use it as feature extractor. Default: True
+        drop_path_rate (float): Stochastic depth rate. Default: 0.
+        layer_scale_init_value (float): Init value for Layer Scale. Default: 1e-6.
+        head_init_scale (float): Init scaling value for classifier weights and biases. Default: 1.
+    """
+
+    def __init__(self, 
+                 depths=[3, 3, 9, 3], dims=[96, 192, 384, 768],
+                 drop_path_rate=0., layer_scale_init_value=1e-6, head_init_scale=1.,
+                 nr_of_classes=7):
+        super().__init__()
+        self.downsample_layers = []  # stem and 3 intermediate downsampling conv layers
+        stem = tf.keras.Sequential([
+            Conv2DTranspose(dims[0], kernel_size=4, strides=4, padding='same'),
+            LayerNormalization(epsilon=1e-6)]
+        )
+        self.downsample_layers.append(stem)
+        for i in range(3):
+            downsample_layer = tf.keras.Sequential([
+                LayerNormalization(epsilon=1e-6),
+                Conv2DTranspose(dims[i+1], kernel_size=2, strides=2, padding='same')]
+            )
+            self.downsample_layers.append(downsample_layer)
+
+        self.stages = []  # 4 feature resolution stages, each consisting of multiple residual blocks
+        dp_rates = [x for x in np.linspace(0, drop_path_rate, sum(depths))]
+        cur = 0
+        for i in range(4):
+            stage = tf.keras.Sequential(
+                [Block(dim=dims[i], drop_path=dp_rates[cur + j],
+                       layer_scale_init_value=layer_scale_init_value, prefix=f'block{i}') for j in range(depths[i])]
+            )
+            self.stages.append(stage)
+            cur += depths[i]
+
+        self.last = Dense(nr_of_classes, activation='softmax')
+
+    def call(self, x):
+        for i in reversed(range(4)):
+            x = self.downsample_layers[i](x)
+            x = self.stages[i](x)
+        x = self.last(x)
+        return x
+
+
+class ConvNeXtDiscriminator(tf.keras.Model):
+    r""" ConvNeXt
+        A Tensorflow keras impl of : `A ConvNet for the 2020s`  -
+          https://arxiv.org/pdf/2201.03545.pdf
+
+    Args:
+        num_classes (int): Number of classes for classification head. Default: 1000
+        depths (tuple(int)): Number of blocks at each stage. Default: [3, 3, 9, 3]
+        dims (int): Feature dimension at each stage. Default: [96, 192, 384, 768]
+        include_top (bool): whether to add head or just use it as feature extractor. Default: True
+        drop_path_rate (float): Stochastic depth rate. Default: 0.
+        layer_scale_init_value (float): Init value for Layer Scale. Default: 1e-6.
+        head_init_scale (float): Init scaling value for classifier weights and biases. Default: 1.
+    """
+
+    def __init__(self, num_classes=1000,
+                 depths=[3, 3, 9, 3], dims=[96, 192, 384, 768], include_top=True,
+                 drop_path_rate=0., layer_scale_init_value=1e-6, head_init_scale=1.,
+                 ):
+        super().__init__()
+        self.downsample_layers = []  # stem and 3 intermediate downsampling conv layers
+        stem = tf.keras.Sequential([
+            Conv2D(dims[0], kernel_size=4, strides=4, padding='same'),
+            LayerNormalization(epsilon=1e-6)]
+        )
+        self.downsample_layers.append(stem)
         for i in range(3):
             downsample_layer = tf.keras.Sequential([
                 LayerNormalization(epsilon=1e-6),
@@ -127,28 +420,21 @@ class ConvNeXt(tf.keras.Model):
             self.stages.append(stage)
             cur += depths[i]
 
-        if self.include_top:
+        self.flatten = Flatten()
-            self.avg = GlobalAveragePooling2D()
+        self.last = Dense(1)
-            self.norm = LayerNormalization(epsilon=1e-6)  # final norm layer
-            self.head = Dense(num_classes)
-        else:
-            self.avg = None
-            self.norm = None
-            self.head = None
 
-    def forward_features(self, x):
+
+    def call(self, x):
+        x = tf.transpose(x, [0,2,3,1])
         for i in range(4):
             x = self.downsample_layers[i](x)
             x = self.stages[i](x)
+        x = self.flatten(x)
+        x = self.last(x)
         return x
 
-    def call(self, x):
+
-        x = self.forward_features(x)
+
-        if self.include_top:
-            x = self.avg(x)
-            x = self.norm(x)
-            x = self.head(x)
-        return x
 
 
 model_urls = {

convunext.py

@@ -0,0 +1,487 @@
+import os
+import time
+import argparse
+import wandb
+import random
+import numpy as np
+import tensorflow as tf
+import scipy.interpolate as spi
+from custom_utils.custom_utils import UV2Mesh, maskFacesBatch
+from tensorboardX import SummaryWriter
+import plotly.graph_objects as go
+import plotly.express as ex
+from plotly.subplots import make_subplots
+
+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 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 convnext.models.convnext_tf import ConvNeXt, ConvNeXtReverse, Block, ConvNeXtReverseWithSkip, ConvNeXtWithSkip
+from stylegan2_ae.generator import Generator
+from custom_utils.graph_utils import Reconstructor_V4_Decoder as GNNDecoder
+from custom_utils.graph_utils import Reconstructor_v2 as SimpleReconstructor
+from dataset import DatasetFactory, StaticDataset
+from custom_utils.custom_utils import batch_to_dict
+from custom_utils.smpl import SMPL
+
+name='convunext_7'
+os.environ["CUDA_VISIBLE_DEVICES"] = "6"
+run = wandb.init(config=tf.compat.v1.flags.FLAGS, sync_tensorboard=True, project='stylegan', name=name,
+         notes='basic pyramid. more balanced levels. take ssim (not ms-ssim) loss at each level. higher lr. on on jumpsuits. \
+uv maps filled with nearest enighbour, up/downscaling with lanczos. only take loss at final level')
+artifact = wandb.Artifact('code', type='code')
+artifact.add_file('auto_encoder.py')
+artifact.add_file('stylegan2_ae/layers/mapping_block.py')
+artifact.add_file('convnext/models/convnext_tf.py')
+run.log_artifact(artifact)
+
+uv2mesh = UV2Mesh(res=256)
+
+
+# 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
+
+
+def plot_3d(mesh_gt, mesh_pred_unposed, faces, collision, collision_faces, collision_color):
+    mesh_pred_unposed[:,1] += 1
+    collision[:,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_unposed[:,0],
+            y=mesh_pred_unposed[:,1],
+            z=mesh_pred_unposed[:,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=collision[:,0],
+            y=collision[:,1],
+            z=collision[:,2],
+            # i, j and k give the vertices of triangles
+            i = collision_faces[:,0],
+            j = collision_faces[:,1],
+            k = collision_faces[:,2],
+            vertexcolor=collision_color
+        )
+    ])
+
+    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
+
+
+def plot_3d_levels(mesh_levels, faces):
+    levels = len(mesh_levels)
+    plot_levels = mesh_levels.copy() #[mesh_levels[i][0] for i in range(levels)]
+    for i in range(levels):
+        plot_levels[i][:,1] = plot_levels[i][:,1] + i * 2
+    meshes = [go.Mesh3d(
+            # 8 vertices of a cube
+            x=mesh[:,0],
+            y=mesh[:,1],
+            z=mesh[:,2],
+            # i, j and k give the vertices of triangles
+            i = faces[:,0],
+            j = faces[:,1],
+            k = faces[:,2]
+        ) for mesh in mesh_levels]
+    fig = go.Figure(data= meshes)
+
+    fig.update_layout(scene_aspectmode='manual',
+                    scene_aspectratio=dict(x=1, y=levels, z=1))
+    fig.update_layout(
+        scene = dict(
+            xaxis = dict(nticks=4, range=[-1,1],),
+                        yaxis = dict(nticks=4, range=[-1,levels*2],),
+                        zaxis = dict(nticks=4, range=[-1,1],),)
+                        )
+
+    return fig
+
+
+class Bottleneck(tf.keras.Model):
+    
+    def __init__(self, dim):
+        super(Bottleneck, self).__init__()
+
+        self.bottleneck = tf.keras.Sequential(
+                [Block(dim=dim, drop_path=0,
+                       layer_scale_init_value=1e-6, prefix=f'bottleneck{j}') for j in range(5)]
+            )
+
+    def call(self, inputs, training=None):
+        x = inputs
+
+        x = self.bottleneck(x, training=training)
+        
+        return x
+
+
+class PyramidLevel(tf.keras.Model):
+
+    def __init__(self, depths=[3, 3, 9, 3], dims=[64, 128, 256, 256]):
+        super(PyramidLevel, self).__init__()
+
+        # reverse_dims = list(reversed(dims[:-1] + [16]))
+        # reverse_dims = dims[:-1] + [16]
+        reverse_dims = [16] + dims[:-1]
+        self.encoder = ConvNeXtWithSkip(depths=depths, dims=dims)
+
+        self.bottleneck = Bottleneck(dims[-1])
+
+        self.decoder = ConvNeXtReverseWithSkip(depths=depths, dims=reverse_dims)
+
+    def call(self, inputs, training=None):
+        x = inputs
+        x, skip_features = self.encoder(x, training=training)
+
+        x = self.bottleneck(x, training=training)
+
+        x = self.decoder([x, skip_features], training=training)
+
+        return x
+
+
+
+def downscale(img, uv_mask=None):
+
+    # img = tf.transpose(tf.squeeze(batch['uv_posed']), [1,2,0]).numpy()
+    # coords = np.transpose(np.stack(np.nonzero((1 - batch['uv_mask'][0]).numpy().astype(bool))), (1,0))
+    # img_idx = np.transpose(np.stack(np.nonzero((batch['uv_mask'][0]).numpy().astype(bool))), (1,0))
+    # img_data = img[img_idx[:,0], img_idx[:,1], :]
+
+    # interpolator =  spi.NearestNDInterpolator(x=img_idx, y=img_data)
+    # img_2 = interpolator(coords)
+    # res = np.zeros((512,256,3))
+    # res[coords[:,0], coords[:,1], :] = img_2
+    # res[img_idx[:,0], img_idx[:,1], :] = img_data
+
+
+    results = [img]
+    if uv_mask is not None:
+        masks = [uv_mask]
+    for i in range(4):
+        # print(np.array(img.shape[:2]) / (2 ** (i + 1)))
+        results.append(tf.image.resize(img, size=np.array(img.shape[1:3]) / (2 ** (i + 1)), method=tf.image.ResizeMethod.LANCZOS5, antialias=True, preserve_aspect_ratio=True))
+        if uv_mask is not None:
+            masks.append(tf.image.resize(uv_mask[:,:,:,tf.newaxis], size=np.array(uv_mask.shape[1:]) / (2 ** (i + 1)), method=tf.image.ResizeMethod.NEAREST_NEIGHBOR, antialias=False)[:,:,:,0])
+    if uv_mask is not None:
+        return results, masks
+    return results
+
+def upscale(imgs):
+    results = []
+    for img in imgs:
+        results.append(tf.image.resize(img, size=(512, 256), method=tf.image.ResizeMethod.LANCZOS5))
+    return results
+
+
+class Pyramid(tf.keras.Model):
+    
+    def __init__(self):
+        super(Pyramid, self).__init__()
+
+        self.level_1 = PyramidLevel(depths=[27], dims=[512])
+        self.level_2 = PyramidLevel(depths=[9, 27], dims=[256, 256])
+        self.level_3 = PyramidLevel(depths=[9, 9, 27], dims=[64, 128, 256])
+        self.level_4 = PyramidLevel(depths=[3, 9, 27, 3], dims=[32, 64, 128, 256])
+        self.level_5 = PyramidLevel(depths=[3, 3, 9, 27, 9], dims=[32, 64, 128, 256, 256])
+
+        self.levels = [self.level_5, self.level_4, self.level_3, self.level_2, self.level_1]
+
+
+    def call(self, inputs, training=None):
+        x, uv_mask, uv_shaped_body = inputs
+
+        x = tf.transpose(x, [0,2,3,1])
+        x_levels = downscale(x)
+
+        output = self.levels[-1](x_levels[-1])
+        output_levels = [output]
+
+        for i in reversed(range(len(x_levels)-1)):
+            level_out = self.levels[i](x_levels[i])
+            output = tf.image.resize(output, size=(output.shape[1] * 2, output.shape[2] * 2), method=tf.image.ResizeMethod.LANCZOS5) + level_out
+            output_levels.append(output)
+
+        x = output
+
+        x = tf.transpose(x, [0,3,1,2])
+        # x = x * uv_mask[:,tf.newaxis,:,:]
+
+        return x, output_levels
+
+class AutoEncoder(tf.keras.Model):
+    
+    def __init__(self):
+        super(AutoEncoder, self).__init__()
+
+
+        # bottleneck_size = 256
+
+        # self.decoder = ConvNeXtReverseWithSkip(depths=list(reversed([3, 3, 9, 3])), dims=list(reversed([bottleneck_size, 128, 64, 16])))
+
+        # self.bottleneck = Bottleneck(bottleneck_size)
+
+        # self.encoder = ConvNeXtWithSkip(depths=[3, 3, 9, 3], dims=[64, 128, 256, bottleneck_size])
+
+        self.level = PyramidLevel(depths=[3, 3, 3, 3], dims=[64, 128])
+
+    def call(self, inputs, training=None):
+        x, uv_mask, uv_shaped_body = inputs
+
+        x = tf.transpose(x, [0,2,3,1])
+        x = tf.image.resize(x, size=(256, 128))
+        # x, skip_features = self.encoder(x, training=training)
+
+        # x = self.bottleneck(x, training=training)
+
+        # x = self.decoder([x, skip_features], training=training)
+
+        x = self.level(x)
+        x = x + tf.transpose(uv_shaped_body, [0,2,3,1])
+        x = tf.transpose(x, [0,3,1,2])
+
+        # x = x * uv_mask[:,:, tf.newaxis]
+        # x = x * uv_mask[:,tf.newaxis,:,:]
+
+        return x
+
+
+@tf.function
+def train_step(batch):
+
+    with tf.GradientTape() as tape:
+
+        pred, levels_pred = model([tf.concat((batch['uv_static'], batch['uv_body_posed']), axis=1), batch['uv_mask'], batch['uv_shaped_body']])
+        # pred = (pred + zero_body) * batch['mesh_mask'][:,:,tf.newaxis]
+
+        mesh_pred = uv2mesh(pred)
+        # mesh_pred = pred
+        mesh_orig = batch['mesh_unposed']
+
+        # uv_unposed_levels, uv_mask_levels = downscale(tf.transpose(batch['uv_unposed'], [0,2,3,1]), batch['uv_mask'])
+        ms_ssim_levels = 0 #tf.reduce_mean([tf.reduce_mean(tf.image.ssim((levels_pred[i] + 1) * uv_mask_levels[len(levels_pred) - i - 1][:,:,:,tf.newaxis], (uv_unposed_levels[len(levels_pred) - i - 1] + 1) * uv_mask_levels[len(levels_pred) - i - 1][:,:,:,tf.newaxis], max_val=10)) for i in range(len(levels_pred))])
+
+        e_loss = 1e-2 * tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']))
+        c_loss = 1e-2 * tf.reduce_mean(collision_penalty(mesh_pred, batch['mesh_body_posed'], batch['mesh_mask']))
+        l1_loss = tf.reduce_mean(tf.reduce_sum((tf.math.abs(batch['uv_unposed'] - pred)), axis=1) * batch['uv_mask'])
+        ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale((tf.transpose(pred, [0,2,3,1]) + 1) * batch['uv_mask'][:,:,:,tf.newaxis], (tf.transpose(batch['uv_unposed'], [0,2,3,1]) + 1) * batch['uv_mask'][:,:,:,tf.newaxis], max_val=10))
+        fft_loss = tf.reduce_mean(tf.reduce_sum(tf.math.abs(tf.signal.fft3d(tf.complex(pred * batch['uv_mask'][:,tf.newaxis,:,:], tf.zeros_like(pred))) - tf.signal.fft3d(tf.complex(batch['uv_unposed'] * batch['uv_mask'][:,tf.newaxis,:,:], tf.zeros_like(batch['uv_unposed'])))), axis=[1,2,3]))
+        # loss = fft_loss_pc # + e_loss + c_loss
+        loss = 1e3 * (1 - ms_ssim)
+
+    mesh_orig_2 = uv2mesh(batch['uv_unposed'])
+    # mesh_skinned = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred)
+    metric = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
+    metric_2 = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig_2 - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
+    euclid_posed = 0 # tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_skinned - batch['mesh_posed']) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
+    gradients = tape.gradient(loss, model.trainable_variables)
+    optimizer.apply_gradients(zip(gradients, model.trainable_variables))
+    return { 'loss': loss, 'edge_loss': e_loss, 'collision_loss': c_loss, 'l1_loss': l1_loss, 'euclid_error': metric, 'ms_ssim': ms_ssim,
+             'fft_loss': fft_loss, 'euclid_error_reconst': metric_2, 
+             'euclid_error_posed': euclid_posed, 'ms_ssim_levels': ms_ssim_levels }
+        
+
+
+@tf.function
+def valid_step(batch):
+    pred, levels_pred = model([tf.concat((batch['uv_static'], batch['uv_body_posed']), axis=1), batch['uv_mask'], batch['uv_shaped_body']])
+
+    mesh_pred = uv2mesh(pred)
+    mesh_orig = batch['mesh_unposed']
+    # uv_unposed_levels, uv_mask_levels = downscale(tf.transpose(batch['uv_unposed'], [0,2,3,1]), batch['uv_mask'])
+    ms_ssim_levels = 0#tf.reduce_mean([tf.reduce_mean(tf.image.ssim(levels_pred[i] + 1, uv_unposed_levels[len(levels_pred) - i - 1] + 1, max_val=10)) for i in range(len(levels_pred))])
+    e_loss = 1e-2 * tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']))
+    c_loss = 1e-2 * tf.reduce_mean(collision_penalty(mesh_pred, batch['mesh_body_posed'], batch['mesh_mask']))
+    l1_loss = tf.reduce_mean(tf.reduce_sum((tf.math.abs(batch['uv_unposed'] - pred)), axis=1) * batch['uv_mask'])
+    ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale((tf.transpose(pred, [0,2,3,1]) + 1) * batch['uv_mask'][:,:,:,tf.newaxis], (tf.transpose(batch['uv_unposed'], [0,2,3,1]) + 1) * batch['uv_mask'][:,:,:,tf.newaxis], max_val=10))
+    fft_loss = tf.reduce_mean(tf.reduce_sum(tf.math.abs(tf.signal.fft3d(tf.complex(pred * batch['uv_mask'][:,tf.newaxis,:,:], tf.zeros_like(pred))) - tf.signal.fft3d(tf.complex(batch['uv_unposed'] * batch['uv_mask'][:,tf.newaxis,:,:], tf.zeros_like(batch['uv_unposed'])))), axis=[1,2,3]))
+    loss = 1e3 * (1 - ms_ssim)
+    mesh_orig_2 = uv2mesh(batch['uv_unposed'])
+    metric = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
+    metric_2 = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig_2 - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
+  
+    return { 'loss': loss, 'edge_loss': e_loss, 'collision_loss': c_loss, 'l1_loss': l1_loss, 'euclid_error': metric, 'ms_ssim': ms_ssim,
+             'fft_loss': fft_loss, 'euclid_error_reconst': metric_2, 'ms_ssim_levels': ms_ssim_levels
+            }
+
+
+def plot_everything(writer, batch: dict, type: str):
+    pred, levels = model([tf.concat((batch['uv_static'], batch['uv_body_posed']), axis=1), batch['uv_mask'], batch['uv_shaped_body']])
+    
+    mesh_pred = uv2mesh(pred)
+    mesh_orig = batch['mesh_unposed']
+    
+    orig_img = tf.transpose(batch['uv_unposed'], [0,2,3,1])
+    pred_img = tf.transpose(pred, [0,2,3,1])
+    img = tf.concat((orig_img, pred_img), axis=2)
+    levels = upscale(levels)
+    mesh_pred_levels = [uv2mesh(tf.transpose(levels[i], [0,3,1,2])) for i in range(len(levels))]
+    with writer.as_default():
+        tf.summary.image('images', img, step=step, max_outputs=4)
+        tf.summary.image('levels', tf.concat(levels, axis=2)  , step=step, max_outputs=4)
+
+    F = np.load("custom_utils/smpl_faces.npy")
+    mask = batch['mesh_mask'].numpy().astype(bool)
+    garment_unposed, faces = maskFacesBatch(mesh_pred.numpy(), F, mask)
+    gt = [mesh_orig[i].numpy()[mask[i], :] for i in range(pred.shape[0])]
+    
+    body = batch['mesh_body_posed'].numpy()
+    mesh_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred).numpy()
+    mesh_posed_levels = [smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_levels[i]).numpy() for i in range(len(levels))]
+    mesh_posed_levels = [mesh_posed_levels[j][0, mask[0], :] for j in range(len(levels))]
+    garment_posed = [mesh_posed[i, mask[i], :] for i in range(body.shape[0])]
+    collision = [np.concatenate((body[i], garment_posed[i]), axis=0) for i in range(len(garment_posed))]
+    collision_faces = [np.concatenate((F, faces[i] + 14475), axis=0) for i in range(len(garment_posed))]
+    collision_color = [np.concatenate((np.ones_like(body[i], dtype=np.int32)*100, np.ones_like(garment_posed[i], dtype=np.int32)*230), axis=0) for i in range(len(garment_posed))]
+
+
+
+    fig_levels = plot_3d_levels(mesh_posed_levels, faces[0])
+
+    figs = { type + "/garment_" + str(i): plot_3d(gt[i], garment_unposed[i], faces[i], collision[i], collision_faces[i], collision_color[i]) for i in range(pred.shape[0]) }
+        # fig = plot_3d(gt[i], garment_unposed[i], faces[i], collision[i], collision_faces[i], collision_color[i])
+    wandb.log( { **figs, **{ type + '/levels_garments': fig_levels } } )
+
+
+if __name__ == '__main__':
+    allow_memory_growth()
+    parser = argparse.ArgumentParser(description='')
+    parser.add_argument('--name', type=str)
+    args = vars(parser.parse_args())
+    
+    smpl = SMPL()
+    # R = smpl.run_rodrigues((np.zeros((8,72), dtype=np.float32)))
+    # zero_body = smpl.run(R, np.zeros((8, 11), dtype=np.float32))[0]
+
+
+    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_train = factory.get_dataset(folder='/home/workspace/local-data/jumpsuit_20/train',
+                                  batch_size=8,
+                                  signature=signature_dynamic)
+    static_dataset_train = StaticDataset(folder='/home/workspace/local-data/jumpsuit_20/train',
+                                   signature=signature_static)
+    dataset_valid = factory.get_dataset(folder='/home/workspace/local-data/jumpsuit_20/valid',
+                                  batch_size=8,
+                                  signature=signature_dynamic)
+    static_dataset_valid = StaticDataset(folder='/home/workspace/local-data/jumpsuit_20/valid',
+                                   signature=signature_static)
+
+
+    # model = AutoEncoder()
+    model = Pyramid()
+    # optimizer = tf.optimizers.Adam(learning_rate=0.002, beta_1=0.0, beta_2=0.99, epsilon=1e-8)
+    lr_schedule = tf.keras.optimizers.schedules.ExponentialDecay(0.001, decay_steps=30000, decay_rate=0.1, staircase=False)
+    optimizer = tf.optimizers.Adam(learning_rate=lr_schedule)
+    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')
+    valid_iterator = iter(dataset_valid)
+
+    step = 0
+    for batch in dataset_train:
+        # batch = batch_to_dict(batch, signature_dynamic)
+        batch = static_dataset_train.expand_batch(batch_to_dict(batch, signature_dynamic))
+
+        train_loss = train_step(batch)
+
+        for key in train_loss.keys():
+            with train_summary_writer.as_default():
+                tf.summary.scalar(key, train_loss[key], step=step)
+
+        # data_valid = batch
+        data_valid = valid_iterator.get_next()
+        data_valid = static_dataset_valid.expand_batch(batch_to_dict(data_valid, signature_dynamic))
+        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)
+
+        if step % 500 == 0:
+            plot_everything(train_summary_writer, batch, 'train')
+            plot_everything(valid_summary_writer, data_valid, 'valid')
+
+            # fig = make_subplots(rows=1, cols=5, start_cell="top-left")
+
+            # for i, l in enumerate(levels):
+            #     fig.add_trace(ex.imshow((l[0] + 1) * 128).data[0], row=1, col=i+1)
+
+                
+
+        print(step)
+
+        step += 1

custom_utils/custom_utils.py

@@ -66,6 +66,33 @@ class UV2Mesh():
         return self.uv_to_mesh(uv)
         
 
+class UV2MeshLabel():
+    def __init__(self, res=256):
+        super(UV2MeshLabel, self).__init__()
+        self.res = res
+        uv_map = np.load('custom_utils/stretched_uv_map_3.npy')
+        self.uv_map_discreet = np.floor(uv_map * self.res).astype(np.int32)
+        m2uv = np.load('custom_utils/m2uv.npy', allow_pickle=True).item()
+        m2uv = m2uv_array(m2uv)
+        self.uv_pixels = self.uv_map_discreet[m2uv]
+
+    # input = channels_first
+    # @tf.function
+    def uv_to_mesh(self, uv):
+        batch_size = uv.shape[0]
+
+        x1_ind = tf.tile(self.uv_pixels[:, 0, :][tf.newaxis, :, :], [batch_size, 1, 1])
+        x2_ind = tf.tile(self.uv_pixels[:, 1, :][tf.newaxis, :, :], [batch_size, 1, 1])
+
+        x1 = tf.gather_nd(uv, x1_ind, batch_dims=1)
+        x2 = tf.gather_nd(uv, x2_ind, batch_dims=1)
+
+        mesh = (x1 + x2) / 2
+        return mesh
+
+    def __call__(self, uv):
+        return self.uv_to_mesh(uv)
+
 
 def _maskFaces(V, F, mask):
     V_mask = mask
@@ -79,6 +106,21 @@ def _maskFaces(V, F, mask):
 
     return V_new, F_new
 
+
+# def _maskFaces_tf(V, F, mask):
+#     V_mask = mask
+#     F_mask = tf.gather(V_mask, F.reshape(-1), axis=0)
+#     F_mask = F_mask.reshape((-1, 3)).sum(axis=1) == 3
+#     F_mask = tf.reduce_sum(tf.reshape(F_mask, (-1, 3)), axis=1) == 3
+
+#     V_ind = tf.zeros(V_mask.shape)
+#     V_ind[V_mask] = tf.range(V_mask.sum())
+#     F_new = tf.gather(V_ind, F[F_mask, :].reshape(-1), axis=0).reshape((-1, 3))
+#     V_new = V[V_mask, :]
+
+#     return V_new, F_new
+
+
 def maskFacesBatch(V, F, mask):
     Vs = []
     Fs = []
@@ -148,3 +190,10 @@ 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
+
+
+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)

custom_utils/graph_utils.py

@@ -1,10 +1,14 @@
+import os
+import sys
 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
+
+sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
+# 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
@@ -31,6 +35,14 @@ def repeat_laplacian(L, times):
     return sp.sparse.coo_matrix((data, (np.squeeze(np.asarray(idx[0,:])), np.squeeze(np.asarray(idx[1,:])))), shape)
 
 
+def convert_pooling(pooling, batchSize):
+    pool_map = np.zeros((np.max(pooling)+1, pooling.shape[0]), dtype=np.float32)
+    for i in range(pooling.shape[0]):
+        pool_map[pooling[i], i] = 1.0
+    pool_map = np.tile(pool_map[np.newaxis, :, :], [batchSize, 1, 1])
+    pool_map = tf.sparse.from_dense(pool_map)
+    return pool_map
+
 
 class GraphConvolution(tf.keras.layers.Layer):
     def __init__(self, shape, activation=None, name='gcn', L=None, batchSize=None, normalization=None):
@@ -174,8 +186,76 @@ class GNNDiscriminator(tf.keras.Model):
         return x
 
 
+class Reconstructor_V4_Decoder(tf.keras.Model):
+    def __init__(self, batchSize, isSkirt=False, n_feature=64):
+        super(Reconstructor_V4_Decoder, self).__init__()
 
 
+        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)]
+        poolings = [convert_pooling(poolings[i][0,:], batchSize) for i in range(len(poolings))]
+
+        self.custom_layers = []
+        for i in range(4,1,-1):
+            self.custom_layers.append(GraphUnPool(pool_map=poolings[i]))
+            self.custom_layers.append(GraphConvolution((int(n_feature/2**(4-(i if i == 4 else i+1))), int(n_feature/2**(4-i)) ), activation=relu, L=laplacians[i], batchSize=batchSize, normalization=BatchNormalization()))
+            for _ 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.pool2 = GraphUnPool(pool_map=poolings[1])
+        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.pool1 = GraphUnPool(pool_map=poolings[0])
+        self.conv1 = GraphConvolution((int(n_feature/2**2), int(n_feature/2**3)), activation=relu, L=laplacians[0], batchSize=batchSize, normalization=BatchNormalization(), name='gcn_1')
+
+        self.conv0 = GraphConvolution((int(n_feature/2**3), 3), L=laplacians[0], batchSize=batchSize, name='gcn_0')
+
+    def call(self, inputs, training=None):
+        x = inputs
+
+        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.pool2(x)
+        x = self.conv2(x, training)
+
+        x = self.pool1(x)
+        x = self.conv1(x, training)
+
+        x = self.conv0(x, training)
+
+        return x
+
+
+
+class Reconstructor_v2(tf.keras.Model):
+    def __init__(self, batchSize, isSkirt=False):
+        super(Reconstructor_v2, self).__init__()
+
+        self.flatten = tf.keras.layers.Flatten()
+
+        self.fc2 = tf.keras.layers.Dense(14475*3, kernel_regularizer=tf.keras.regularizers.l2)
+
+        self.reshape2 = tf.keras.layers.Reshape((14475, 3))
+        # self.drop = Dropout(0.2)
+        # self.fc3 = Dense(512)
+        self.fc1 = tf.keras.layers.Dense(3)
+
+    def call(self, inputs, training=None):
+        x = self.flatten(inputs)
+        x = self.fc2(x)
+        # x = self.drop(x, training=training)
+        x = self.reshape2(x)
+        # x = self.fc3(x)
+        x = self.fc1(x)
+        return x
 
 
 class DatasetTest(tf.data.Dataset):
@@ -195,13 +275,16 @@ class DatasetTest(tf.data.Dataset):
 
 if __name__ == '__main__':
     import os
-    os.environ["CUDA_VISIBLE_DEVICES"] = "1"
+    os.environ["CUDA_VISIBLE_DEVICES"] = "2"
 
     batch_size = 4
-    input = tf.ones((batch_size, 14475, 3))
+    # input = tf.ones((batch_size, 14475, 3))
 
-    ge = GNNDiscriminator(batchSize=batch_size, n_feature=32)
+    # ge = GNNDiscriminator(batchSize=batch_size, n_feature=32)
 
-    x = ge(input)
+    # x = ge(input)
+
+    gd = Reconstructor_V4_Decoder(batch_size)
+    x = gd(tf.ones((batch_size, 21, 64)))
 
     print()

cvae.py

@@ -0,0 +1,584 @@
+import os
+os.environ["CUDA_VISIBLE_DEVICES"] = "2"
+import tensorflow as tf
+
+gpus = tf.config.experimental.list_physical_devices('GPU')
+if gpus:
+    try:
+        # Currently, memory growth needs to be the same across GPUs
+        for gpu in gpus:
+            tf.config.experimental.set_memory_growth(gpu, True)
+        logical_gpus = tf.config.experimental.list_logical_devices('GPU')
+        print(len(gpus), "Physical GPUs,", len(logical_gpus), "Logical GPUs")
+    except RuntimeError as e:
+        # Memory growth must be set before GPUs have been initialized
+        print(e)
+
+import time
+import argparse
+import wandb
+import random
+import numpy as np
+import scipy as sp
+from custom_utils.custom_utils import UV2Mesh, maskFacesBatch, _maskFaces
+from tensorboardX import SummaryWriter
+import plotly.graph_objects as go
+import plotly.express as ex
+
+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 losses import d_logistic, d_logistic_r1_reg, g_logistic_non_saturating, g_logistic_ns_pathreg, edge_loss, uv_mask_loss, bend_loss
+# from losses import collision_penalty as collision_penalty
+from losses import collision_loss as collision_penalty
+from convnext.models.convnext_tf import ConvNeXt, ConvNeXtReverse, Block
+from stylegan2_ae.generator import Generator
+from custom_utils.graph_utils import Reconstructor_V4_Decoder as GNNDecoder, sparse_to_tensor
+from custom_utils.graph_utils import Reconstructor_v2 as SimpleReconstructor
+from dataset import DatasetFactory, StaticDataset
+from custom_utils.custom_utils import batch_to_dict
+from custom_utils.smpl import SMPL
+from auto_encoder_body import AutoEncoder as BodyAE
+from auto_encoder_curv import AutoEncoder as CurvatureAE
+from auto_encoder_static import AutoEncoder as StaticAE
+
+
+name='cvae_10_2'
+def init_wandb():
+    run = wandb.init(project='vae', name=name,
+            notes='learn offsets. smaller flattened latent vectors. with 25-85k pretraining. smaller batchsize. fixed balancing of reconst and kl loss for baseline.\
+new implementation for kl loss. no lr schedule. 25 jumpsuits, average reconst_loss before calculation, linearly increasing balancing, up to 1000 from 250.\
+offsets over average unposed garment. gamma balancing. remove global orientation (in posed body input and in pose_shape parameters. larger batchsize\
+add edge and collision (fixed) loss. rebalanced l1 loss. increasing edge and coll coeff. retrained body encoder with deoriented bodies \
+don\'t freeze encoders, lower lr')
+    artifact = wandb.Artifact('code', type='code')
+    artifact.add_file('cvae.py')
+    artifact.add_file('dataset.py')
+    artifact.add_file('losses.py')
+    artifact.add_file('convnext/models/convnext_tf.py')
+    run.log_artifact(artifact)
+
+uv2mesh = UV2Mesh(res=256)
+
+def plot_3d(mesh_gt, mesh_pred_unposed, faces, collision, collision_faces, collision_color):
+    mesh_pred_unposed[:,1] += 1
+    collision[:,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_unposed[:,0],
+            y=mesh_pred_unposed[:,1],
+            z=mesh_pred_unposed[:,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=collision[:,0],
+            y=collision[:,1],
+            z=collision[:,2],
+            # i, j and k give the vertices of triangles
+            i = collision_faces[:,0],
+            j = collision_faces[:,1],
+            k = collision_faces[:,2],
+            vertexcolor=collision_color
+        )
+    ])
+
+    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 Bottleneck(tf.keras.Model):
+    
+    def __init__(self, dim):
+        super(Bottleneck, self).__init__()
+
+        self.bottleneck = tf.keras.Sequential(
+                [Block(dim=dim, drop_path=0,
+                       layer_scale_init_value=1e-6, prefix=f'bottleneck{j}') for j in range(5)]
+            )
+
+    def call(self, inputs, training=None):
+        x = inputs
+
+        x = self.bottleneck(x, training=training)
+        
+        return x
+
+
+class Encoder(tf.keras.Model):
+    
+    def __init__(self, latent_dim):
+        super(Encoder, self).__init__()
+
+        self.encoder = ConvNeXt(depths=[3, 3, 9, 3], dims=[32, 64, 128, 256], include_top=False)
+        self.flatten = tf.keras.layers.Flatten()
+        self.mean = tf.keras.layers.Dense(latent_dim)
+        self.log_var = tf.keras.layers.Dense(latent_dim)
+
+    def call(self, inputs, training=None):
+        x = inputs
+
+        x = self.encoder(x)
+        x = self.flatten(x)
+        mean = self.mean(x)
+        log_var = self.log_var(x)
+
+        return mean, log_var
+
+
+class Decoder(tf.keras.Model):
+    
+    def __init__(self):
+        super(Decoder, self).__init__()
+
+        self.dense = tf.keras.layers.Dense(16 * 8 * 80)
+        self.reshape = tf.keras.layers.Reshape((16, 8, 80))
+        self.decoder = ConvNeXtReverse(depths=list(reversed([3, 3, 9, 3])), dims=list(reversed([256, 128, 64, 16])))
+
+
+    def call(self, inputs, training=None):
+        x = inputs
+
+        x = self.dense(x)
+        x = self.reshape(x)
+        x = self.decoder(x)
+
+        return x
+
+
+class CVAE(tf.keras.Model):
+    
+    def __init__(self, latent_dim=128):
+        super(CVAE, self).__init__()
+        self.latent_dim = latent_dim
+
+        model : StaticAE = StaticAE()
+        load_model(model, 'checkpoints/ae_static_3')
+        self.static_enc = tf.keras.Sequential([model.encoder, model.bottleneck])
+        # self.static_enc.trainable = False
+
+        model : BodyAE = BodyAE()
+        load_model(model, 'checpoints/ae_body_5')
+        self.body_enc = tf.keras.Sequential([model.encoder, model.bottleneck])
+        # self.body_enc.trainable = False
+
+        model : CurvatureAE = CurvatureAE()
+        load_model(model, 'checpoints/ae_curvature_6')
+        self.curv_enc = tf.keras.Sequential([model.encoder, model.bottleneck])
+        # self.curv_enc.trainable = False
+
+        self.encoder = Encoder(latent_dim)
+        self.decoder = Decoder()
+
+
+    # @tf.function
+    def sample(self, uv_template, uv_body_posed, uv_curv, uv_mask, eps=None):
+        if eps is None:
+            eps = tf.random.normal(shape=(4, self.latent_dim))
+
+        uv_template = tf.tile(tf.transpose(uv_template, [1,2,0])[tf.newaxis,:,:,:], [4,1,1,1])
+        uv_body_posed = tf.tile(tf.transpose(uv_body_posed, [1,2,0])[tf.newaxis,:,:,:], [4,1,1,1])
+        uv_curv = tf.tile(uv_curv[tf.newaxis,:,:,:], [4,1,1,1])
+        uv_mask = tf.tile(uv_mask[tf.newaxis,:,:], [4,1,1])
+
+        # uv_template, uv_body_posed = tf.transpose(uv_template, [0,2,3,1]), tf.transpose(uv_body_posed, [0,2,3,1])
+        z_static = self.static_enc(uv_template)
+        z_body = self.body_enc(uv_body_posed)
+        z_curv = self.curv_enc(uv_curv)
+
+        z_cond = tf.concat([z_static, z_body, z_curv], axis=-1)
+        z = tf.concat([eps, z_cond], axis=-1)
+
+        x = self.decoder(z)
+        x = x + avg_jumpsuit[tf.newaxis,:,:,:]
+        x = x * uv_mask[:,:,:,tf.newaxis]
+
+        return x
+
+
+    def reparameterize(self, mean, logvar):
+        eps = tf.random.normal(shape=mean.shape)
+        return eps * tf.exp(logvar * .5) + mean
+
+
+    def decode(self, z, apply_sigmoid=False):
+        logits = self.decoder(z)
+        if apply_sigmoid:
+            probs = tf.sigmoid(logits)
+            return probs
+        return logits
+
+
+    def call(self, inputs, training=None):
+        uv_unposed, uv_template, uv_body_posed, uv_curv, uv_mask, uv_shaped_body = inputs
+        uv_unposed, uv_template, uv_body_posed = tf.transpose(uv_unposed, [0,2,3,1]), tf.transpose(uv_template, [0,2,3,1]), tf.transpose(uv_body_posed, [0,2,3,1])
+        z_static = self.static_enc(uv_template)
+        z_body = self.body_enc(uv_body_posed)
+        z_curv = self.curv_enc(uv_curv)
+
+        x = tf.concat([uv_unposed, uv_template, uv_body_posed, uv_curv], axis=-1)
+
+        mean, logvar = self.encoder(x)
+        z = self.reparameterize(mean, logvar)
+        z = tf.concat([z, z_static, z_body, z_curv], axis=-1)
+        x = self.decoder(z)
+
+
+        # x = tf.transpose(x, [0,3,1,2])
+        # x = x * uv_mask[:,tf.newaxis,:,:]
+        # x = tf.image.resize(x, (512, 256), tf.image.ResizeMethod.LANCZOS5)
+        # x = x + tf.transpose(uv_shaped_body, [0,2,3,1])
+        x = x + avg_jumpsuit[tf.newaxis,:,:,:]
+        x = x * uv_mask[:,:,:,tf.newaxis]
+
+        return x, mean, logvar, z
+
+
+def log_normal_pdf(sample, mean, logvar, raxis=1):
+    log2pi = tf.math.log(2. * np.pi)
+    res = tf.reduce_sum(-.5 * ((sample - mean) ** 2. * tf.exp(-logvar) + logvar + log2pi), axis=raxis)
+    return res
+
+
+def kl_loss(mean, logvar):
+    kl_loss = 1 + logvar - tf.square(mean) - tf.exp(logvar)
+    kl_loss = tf.reduce_mean(kl_loss)
+    kl_loss *= -0.5
+    return kl_loss
+
+
+@tf.function
+def train_step(batch):
+
+    with tf.GradientTape() as tape:
+        uv_curvature = tf.one_hot(batch['uv_curvature'], 7)
+
+        pred, mean, logvar, z = model([batch['uv_unposed'], batch['uv_static'], batch['uv_body_posed'], uv_curvature, batch['uv_mask'], batch['uv_shaped_body']])
+
+        mesh_pred = uv2mesh(tf.transpose(pred, [0,3,1,2]))
+        mesh_pred_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred)
+        mesh_orig = batch['mesh_unposed']
+
+
+        e_loss = tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']))
+        c_loss = tf.reduce_mean(collision_penalty(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask']))
+        l1_loss = tf.reduce_sum(tf.reduce_sum((tf.math.abs(tf.transpose(batch['uv_unposed'], [0,2,3,1]) - pred)), axis=-1) * batch['uv_mask'], axis=[1, 2])
+        # ms_ssim = tf.image.ssim_multiscale(tf.transpose(pred, [0,2,3,1]) + 1, tf.transpose(batch['uv_unposed'], [0,2,3,1]) + 1, max_val=2)
+        ms_ssim = tf.image.ssim_multiscale(pred + 1, tf.transpose(batch['uv_unposed'], [0,2,3,1]) + 1, max_val=3)
+        # ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale((tf.transpose(pred, [0,2,3,1]) + 1) * batch['uv_mask'][:,:,:,tf.newaxis], (tf.transpose(batch['uv_unposed'], [0,2,3,1]) + 1) * batch['uv_mask'][:,:,:,tf.newaxis], max_val=2))
+        rfft_loss = 0 #tf.reduce_mean(tf.reduce_sum(tf.math.abs(tf.signal.rfft3d(pred) - tf.signal.rfft3d(batch['uv_unposed'])), axis=[1,2,3]))
+        fft_loss = tf.reduce_sum(tf.math.abs(tf.signal.fft3d(tf.complex(pred, tf.zeros_like(pred))) - tf.signal.fft3d(tf.complex(tf.transpose(batch['uv_unposed'], [0,2,3,1]), tf.zeros_like(tf.transpose(batch['uv_unposed'], [0,2,3,1]))))), axis=[1,2,3])
+        fft_loss_pc, fft_loss_pc_low, fft_loss_pc_high = 0,0,0 #fft_loss_per_channel(pred, batch['uv_unposed'])
+        tv_loss = tf.image.total_variation(pred)
+
+        reconstr_loss = l1_loss / 4
+        reconstr_loss = tf.reduce_mean(reconstr_loss) + reg_coeff *(e_loss + c_loss)
+        # logpx_z = -reconstr_loss
+        # logpz = log_normal_pdf(z[:,:model.latent_dim], 0., 0.)
+        # logqz_x = log_normal_pdf(z[:,:model.latent_dim], mean, logvar)
+        # loss = -tf.reduce_mean(logpx_z + balancing_coeff * (logpz - logqz_x))
+        kl = kl_loss(mean, logvar)
+        loss = reconstr_loss / (2 * (gamma ** 2)) + balancing_coeff * kl
+        # loss = reconstr_loss + balancing_coeff * kl
+        ms_ssim = tf.reduce_mean(ms_ssim)
+        l1_loss = tf.reduce_mean(l1_loss)
+        fft_loss = tf.reduce_mean(fft_loss)
+        tv_loss = tf.reduce_mean(tv_loss)
+        mean = tf.reduce_mean(mean)
+        logvar = tf.reduce_mean(logvar)
+
+    mesh_orig_2 = uv2mesh(batch['uv_unposed'])
+    metric = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
+    metric_2 = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig_2 - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
+    euclid_posed = 0 # tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_skinned - batch['mesh_posed']) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
+    gradients = tape.gradient(loss, model.trainable_variables)
+    # gradients = [tf.clip_by_norm(g, 1.0) for g in gradients]
+    # gradients = [tf.clip_by_value(g, -1.0, 1.0) for g in gradients]
+    optimizer.apply_gradients(zip(gradients, model.trainable_variables))
+    return { 'train/loss': loss, 'train/edge_loss': e_loss, 'train/collision_loss': c_loss, 'train/l1_loss': l1_loss, 'train/euclid_error': metric, 'train/ms_ssim': ms_ssim,
+             'train/fft_loss': fft_loss, 'train/euclid_error_reconst': metric_2, 'train/rfft_loss': rfft_loss, 'train/fft_oc_low': fft_loss_pc_low, 'train/fft_pc_high': fft_loss_pc_high,
+             'train/fft_loss_per_channel': fft_loss_pc, 'train/euclid_error_posed': euclid_posed, 'train/reconstr_loss': reconstr_loss, 'train/tv_loss': tv_loss,
+             'train/mean': mean, 'train/logvar': logvar, 'train/kl_loss': kl, 'train/reconst_loss_norm': reconstr_loss / (2 * (gamma ** 2)), 'train/var': np.e ** logvar,
+             'train/reg_coeff': reg_coeff }
+        
+
+
+@tf.function
+def valid_step(batch):
+    uv_curvature = tf.one_hot(batch['uv_curvature'], 7)
+    pred, mean, logvar, z = model([batch['uv_unposed'], batch['uv_static'], batch['uv_body_posed'], uv_curvature, batch['uv_mask'], batch['uv_shaped_body']])
+
+    mesh_pred = uv2mesh(tf.transpose(pred, [0,3,1,2]))
+    mesh_orig = batch['mesh_unposed']
+    e_loss = 1e-2 * tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']))
+    c_loss = 1e-2 * tf.reduce_mean(collision_penalty(mesh_pred, batch['mesh_body_posed'], batch['mesh_mask']))
+    l1_loss = tf.reduce_sum(tf.reduce_sum((tf.math.abs(tf.transpose(batch['uv_unposed'], [0,2,3,1]) - pred)), axis=-1) * batch['uv_mask'], axis=[1, 2])
+    b_loss = 0 #1e-3 * bend_loss(mesh_pred, smpl_faces, laplacian_loss, batch['mesh_mask'])
+    # ms_ssim = tf.image.ssim_multiscale(tf.transpose(pred, [0,2,3,1]) + 1, tf.transpose(batch['uv_unposed'], [0,2,3,1]) + 1, max_val=2)
+    ms_ssim = tf.image.ssim_multiscale(pred + 1, tf.transpose(batch['uv_unposed'], [0,2,3,1]) + 1, max_val=3)
+    fft_loss = 0 #tf.reduce_sum(tf.math.abs(tf.signal.fft3d(tf.complex(pred, tf.zeros_like(pred))) - tf.signal.fft3d(tf.complex(tf.transpose(batch['uv_unposed'], [0,2,3,1]), tf.zeros_like(tf.transpose(batch['uv_unposed'], [0,2,3,1]))))), axis=[1,2,3])
+    # mesh_orig_2 = uv2mesh(batch['uv_unposed'])
+    metric = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
+    metric_2 = 0 #tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig_2 - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
+  
+    reconstr_loss = l1_loss
+    # logpx_z = -reconstr_loss
+    # logpz = log_normal_pdf(z[:,:model.latent_dim], 0., 0.)
+    # logqz_x = log_normal_pdf(z[:,:model.latent_dim], mean, logvar)
+    # loss = -tf.reduce_mean(logpx_z + balancing_coeff * (logpz - logqz_x))
+    kl = kl_loss(mean, logvar)
+    # loss = reconstr_loss / (2 * (gamma ** 2)) + kl
+    loss = reconstr_loss + balancing_coeff * kl
+    ms_ssim = tf.reduce_mean(ms_ssim)
+    l1_loss = tf.reduce_mean(l1_loss)
+    fft_loss = tf.reduce_mean(fft_loss)
+    b_loss = tf.reduce_mean(b_loss)
+    reconstr_loss = tf.reduce_mean(reconstr_loss)
+    mean = tf.reduce_mean(mean)
+    logvar = tf.reduce_mean(logvar)
+
+    return { 'valid/loss': loss, 'valid/edge_loss': e_loss, 'valid/collision_loss': c_loss, 'valid/l1_loss': l1_loss, 'valid/euclid_error': metric, 'valid/ms_ssim': ms_ssim,
+             'valid/fft_loss': fft_loss, 'valid/euclid_error_reconst': metric_2, 'train/bend_loss': b_loss, 'train/reconstr_loss': reconstr_loss,
+             'valid/mean': mean, 'valid/logvar': logvar, 'valid/kl_loss': kl, 'valid/var': np.e ** logvar }
+
+
+def plot_everything(batch: dict, type: str):
+    uv_curvature = tf.one_hot(batch['uv_curvature'], 7)
+    pred, mean, logvar, z = model([batch['uv_unposed'], batch['uv_static'], batch['uv_body_posed'], uv_curvature, batch['uv_mask'], batch['uv_shaped_body']])
+    
+    mesh_pred = uv2mesh(tf.transpose(pred, [0,3,1,2]))
+    mesh_orig = batch['mesh_unposed']
+    
+    orig_img = tf.transpose(batch['uv_unposed'], [0,2,3,1])
+    uv_static_img = tf.transpose(batch['uv_static'], [0,2,3,1])
+    uv_body_posed_img = tf.transpose(batch['uv_body_posed'], [0,2,3,1])
+    pred_img = pred
+    img = tf.concat((uv_static_img, uv_body_posed_img, orig_img, pred_img), axis=2)
+    imgs = [wandb.Image((img[i] + 1) / 2.0) for i in range(4)]
+
+    F = np.load("custom_utils/smpl_faces.npy")
+    mask = batch['mesh_mask'].numpy().astype(bool)
+    garment_unposed, faces = maskFacesBatch(mesh_pred.numpy(), F, mask)
+    gt = [mesh_orig[i].numpy()[mask[i], :] for i in range(pred.shape[0])]
+    
+    body = batch['mesh_body_posed'].numpy()
+    mesh_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred).numpy()
+    garment_posed = [mesh_posed[i, mask[i], :] for i in range(body.shape[0])]
+    collision = [np.concatenate((body[i], garment_posed[i]), axis=0) for i in range(len(garment_posed))]
+    collision_faces = [np.concatenate((F, faces[i] + 14475), axis=0) for i in range(len(garment_posed))]
+    collision_color = [np.concatenate((np.ones_like(body[i], dtype=np.int32)*100, np.ones_like(garment_posed[i], dtype=np.int32)*230), axis=0) for i in range(len(garment_posed))]
+
+    figs = { type + "/garment_" + str(i): plot_3d(gt[i], garment_unposed[i], faces[i], collision[i], collision_faces[i], collision_color[i]) for i in range(4) }
+    
+    mean_hist = wandb.Histogram(sequence=mean.numpy().mean(axis=0), num_bins=model.latent_dim)
+    logvar_hist = wandb.Histogram(sequence=logvar.numpy().mean(axis=0), num_bins=model.latent_dim)
+
+    return { **figs, **{ type + '/images': imgs, type + '/mean_hist': mean_hist, type + '/logvar_hist': logvar_hist } }
+
+
+def load_model(model, checkpoint):
+    ckpt = tf.train.Checkpoint(model=model)
+    manager = tf.train.CheckpointManager(ckpt, checkpoint, max_to_keep=1)
+    ckpt.restore(manager.latest_checkpoint)
+
+
+def sample(folder):
+    load_model(model,folder)
+    for batch in dataset_train:
+        batch = static_dataset_train.expand_batch(batch_to_dict(batch, signature_dynamic))
+        uv_curvature = tf.one_hot(batch['uv_curvature'], 7)
+
+        pred = model.sample(batch['uv_static'][0], batch['uv_body_posed'][0], uv_curvature[0], batch['uv_mask'][0])
+        pose, shape = batch['pose_shape'][0, :72], batch['pose_shape'][0, 72:]
+        pose, shape = tf.tile(pose[tf.newaxis,:], [4,1]), tf.tile(shape[tf.newaxis,:], [4,1])
+
+        mesh_pred = uv2mesh(tf.transpose(pred, [0,3,1,2]))
+        mesh_posed = smpl.run_garment(pose, shape, mesh_pred).numpy()
+        mesh_orig = batch['mesh_posed'][0].numpy()
+
+        body = batch['mesh_body_posed'][0].numpy()
+        body_color = np.ones_like(body, dtype=np.int32)*100
+
+        F = np.load("custom_utils/smpl_faces.npy")
+        mask = batch['mesh_mask'].numpy().astype(bool)[0]
+        _, faces = _maskFaces(np.ones((14475, 3)), F, mask)
+
+        batch = 0
+
+        meshes = []
+        meshes.append(go.Mesh3d(
+                    # 8 vertices of a cube
+                    x=mesh_orig[mask,0],
+                    y=mesh_orig[mask,1],
+                    z=mesh_orig[mask,2],
+                    # i, j and k give the vertices of triangles
+                    i = faces[:,0],
+                    j = faces[:,1],
+                    k = faces[:,2]
+                    ))
+
+        for i in range(4):
+            x = mesh_posed[i][mask, :].copy()
+            x[:, 1] += 1 * (i + 1)
+            meshes.append(go.Mesh3d(
+                        # 8 vertices of a cube
+                        x=x[:,0],
+                        y=x[:,1],
+                        z=x[:,2],
+                        # i, j and k give the vertices of triangles
+                        i = faces[:,0],
+                        j = faces[:,1],
+                        k = faces[:,2]
+                        ))
+
+        for i in range(5):
+            x = body.copy()
+            x[:, 1] += 1 * i
+            meshes.append(go.Mesh3d(
+                        # 8 vertices of a cube
+                        x=x[:,0],
+                        y=x[:,1],
+                        z=x[:,2],
+                        # i, j and k give the vertices of triangles
+                        i = F[:,0],
+                        j = F[:,1],
+                        k = F[:,2],
+                        vertexcolor=body_color
+                        ))
+
+
+        fig = go.Figure(data=meshes)
+
+        fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=3, z=1))
+        fig.update_layout(
+            scene = dict(
+                xaxis = dict(nticks=4, range=[-1,1],),
+                            yaxis = dict(nticks=4, range=[-1,5],),
+                            zaxis = dict(nticks=4, range=[-1,1],),)
+                            )
+
+        fig.show()
+
+        print()
+
+if __name__ == '__main__':
+    # init_wandb()
+    # allow_memory_growth()
+    parser = argparse.ArgumentParser(description='')
+    parser.add_argument('--name', type=str)
+    args = vars(parser.parse_args())
+
+    batch_size = 12
+    
+    smpl = SMPL()
+
+    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,),
+                          'uv_curvature': (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_train = factory.get_dataset(folder='/home/workspace/local-data/jumpsuit_20/train', batch_size=batch_size, signature=signature_dynamic)
+    static_dataset_train = StaticDataset(folder='/home/workspace/local-data/jumpsuit_20/train', signature=signature_static)
+    dataset_valid = factory.get_dataset(folder='/home/workspace/local-data/jumpsuit_20/valid', batch_size=batch_size, signature=signature_dynamic)
+    static_dataset_valid = StaticDataset(folder='/home/workspace/local-data/jumpsuit_20/valid', signature=signature_static)
+
+
+    model = CVAE()
+
+
+    # lr_schedule = tf.keras.optimizers.schedules.ExponentialDecay(0.0001, decay_steps=30000, decay_rate=0.1, staircase=False)
+    optimizer = tf.optimizers.Adam(learning_rate=0.0001)
+    balancing_coeff = 250
+    reg_coeff = 1
+    folder = 'models/ae_logs/'
+    valid_iterator = iter(dataset_valid)
+    gamma = 1
+    balancing_activated_step = 0
+    avg_jumpsuit = np.load('custom_data/jumpsuit_20_avg_uv.npy')
+
+
+    sample('checkpoints/cvae_10')
+
+
+    ckpt = tf.train.Checkpoint(model=model)
+    manager = tf.train.CheckpointManager(ckpt, 'checkpoints/' + name, max_to_keep=3)
+
+    step = 0
+    for batch in dataset_train:
+        batch = static_dataset_train.expand_batch(batch_to_dict(batch, signature_dynamic))
+
+        train_loss = train_step(batch)
+
+        # data_valid = valid_iterator.get_next()
+        # data_valid = static_dataset_valid.expand_batch(batch_to_dict(data_valid, signature_dynamic))
+        # valid_loss = valid_step(data_valid)
+        valid_loss = {}
+
+        to_log = { **train_loss, **valid_loss, **{ 'current_step': step} }
+        if step % 2500 == 0:
+            train_figs = plot_everything(batch, 'train')
+            # valid_figs = plot_everything(data_valid, 'valid')
+            valid_figs = {}
+            to_log = { **to_log, **train_figs, **valid_figs }
+
+        if step % 10000 == 0:
+            manager.save(checkpoint_number=step)
+
+        if train_loss['train/euclid_error'] < 0.01:
+            if balancing_activated_step == 0:
+                balancing_activated_step = step
+            balancing_coeff = balancing_coeff + 1000 / 30000
+        if train_loss['train/euclid_error'] < 0.015:
+            reg_coeff = reg_coeff + 100 / 15000
+        # if step > 15000 and balancing_coeff <= 1e-3:
+        #     balancing_coeff = 1e-2
+        # if step > 20000 and balancing_coeff <= 1e-3:
+        #     balancing_coeff = 1e-2
+        # if step > 25000 and balancing_coeff <= 1e-2:
+        #     balancing_coeff = 1e-1
+        # if step > 30000 and balancing_coeff <= 1e-1:
+        #     balancing_coeff = 1
+
+        # gamma = 0.99 * gamma + 0.01 * train_loss['train/reconstr_loss']        
+
+        to_log = { **to_log, **{ 'train/balancing_coeff': balancing_coeff, 'train/gamma': gamma } }
+        
+        wandb.log(to_log, step = step)
+
+        print(step)
+
+        step += 1

dataset.py

@@ -6,7 +6,7 @@ import random
 import numpy as np
 import tensorflow as tf
 
-from custom_utils.custom_utils import nest_conditions
+from custom_utils.custom_utils import nest_conditions, batch_to_dict
 
 
 
@@ -101,7 +101,7 @@ 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 = sorted(os.listdir(uv_static_folder))#[453:503]
         frames = [frame.split('.')[0] for frame in frames] # .split('_')[0]
         random.shuffle(frames)
 
@@ -119,6 +119,8 @@ class DatasetFactory():
         for key in signature.keys():
             if key == 'id':
                 tf_signature.append(tf.TensorSpec(shape = (), dtype = tf.string))
+            elif key == 'uv_curvature' or key == 'mesh_curvature':
+                tf_signature.append(tf.TensorSpec(shape = signature[key], dtype = tf.int32))
             else:
                 tf_signature.append(tf.TensorSpec(shape = signature[key], dtype = tf.float32))
         return tf_signature
@@ -139,16 +141,27 @@ class DatasetFactory():
                 item['id'] = id
             if 'pose_shape' in signature.keys():
                 pose_shape = np.load(os.path.join(folder, 'pose_shape',  id + '_' + frame_nr + '.npy'))
+                
+                ##############################
+                pose_shape[:3] = 0
+                ##############################
+
                 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 '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])
+                uv_posed_body_no_orient = np.load(os.path.join(folder, 'uv_posed_body_no_orient', id + '_' + frame_nr + '.npz')); uv_posed_body_no_orient = uv_posed_body_no_orient['uv_posed_body_no_orient']; uv_posed_body_no_orient = np.transpose(uv_posed_body_no_orient, [2, 0, 1])
-                item['uv_body_posed'] = uv_body_posed
+                item['uv_body_posed'] = uv_posed_body_no_orient
             if 'mesh_body_posed' in signature.keys():
-                mesh_body_posed = np.load(os.path.join(folder, 'mesh_body_posed', id + '_' + frame_nr + '.npy'))
+                mesh_posed_body_no_orient = np.load(os.path.join(folder, 'mesh_posed_body_no_orient', id + '_' + frame_nr + '.npz')); mesh_posed_body_no_orient = mesh_posed_body_no_orient['mesh_posed_body_no_orient']
-                item['mesh_body_posed'] = mesh_body_posed
+                item['mesh_body_posed'] = mesh_posed_body_no_orient
             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
@@ -158,7 +171,12 @@ class DatasetFactory():
             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
-
+            if 'uv_curvature' in signature.keys():
+                uv_curvature = np.load(os.path.join(folder, 'uv_curvature', id + '_' + frame_nr + '.npz')); uv_curvature = uv_curvature['curvature_uv']
+                item['uv_curvature'] = uv_curvature.astype(np.int32)
+            if 'mesh_curvature' in signature.keys():
+                mesh_curvature = np.load(os.path.join(folder, 'mesh_curvature', id + '_' + frame_nr + '.npz')); mesh_curvature = mesh_curvature['mesh_curvature']
+                item['mesh_curvature'] = mesh_curvature.astype(np.int32)
 
             # item = { **static_info, **item }
             as_list = []
@@ -168,6 +186,10 @@ class DatasetFactory():
 
             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(lambda i: tf.py_function(func=read_item, 
                                                inp=[i], 
                                                Tout=self._build_tf_signature(signature)
@@ -186,44 +208,232 @@ class DatasetFactory():
 
 
 
-os.environ["CUDA_VISIBLE_DEVICES"] = "3"
 
 
 if __name__ == '__main__':
-    from custom_utils.custom_utils import batch_to_dict
+    os.environ["CUDA_VISIBLE_DEVICES"] = "6"
+    from custom_utils.custom_utils import batch_to_dict, UV2Mesh, _maskFaces
+    import plotly.graph_objects as go
+    # import plotly.express as ex
+    # import cmapy as cm
+    uv2mesh = UV2Mesh()
+
+    F = np.load("custom_utils/smpl_faces.npy")
+
+
+    # dataset = DatasetIdx('/home/workspace/data/dataset/stylegan_new/train')
+    # for d in dataset:
+    #     print(d)
 
     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_body_posed': (14475, 3,),
-                          'mesh_posed': (14475, 3,),
+                        #   'mesh_posed': (14475, 3,),
-                        #   'mesh_unposed': (14475, 3,),
+                          'mesh_unposed': (14475, 3,),
-                        #   'uv_unposed': (3, 512, 256,),
+                          'uv_unposed': (3, 512, 256,),
+                        #   'uv_curvature': (512, 256),
+                        #   'mesh_curvature':  (14475)
                         }
     signature_static = { #'uv_static': (3, 512, 256,),
                         #  'mesh_static': (14475, 3,),
                          'mesh_mask': (14475,),
-                        #  'uv_mask': (512, 256,),
+                         'uv_mask': (512, 256,),
                         #  'uv_shaped_body': (3, 512, 256,),
                           }
 
-    dataset = factory.get_dataset(folder='/home/workspace/data/dataset/stylegan_new/train',
+    dataset = factory.get_dataset(folder='/home/workspace/local-data/jumpsuit_20/train',
-                                  batch_size=1024,
+                                  batch_size=1,
                                   signature=signature_dynamic)
-    static_dataset = StaticDataset(folder='/home/workspace/data/dataset/stylegan_new/train',
+    static_dataset = StaticDataset(folder='/home/workspace/local-data/jumpsuit_20/train',
                                    signature=signature_static)
 
-    @tf.function
+
+
+
+    # sum_uv = np.zeros((3, 512, 256), dtype=np.float32)
+    # sum_mask = np.zeros((512, 256), dtype=np.float32)
+    # for i, batch in enumerate(dataset):
+    #     print(i)
+    #     batch = static_dataset.expand_batch(batch_to_dict(batch, signature_dynamic))
+    #     sum_uv += batch['uv_unposed'].numpy().sum(axis=0)
+    #     sum_mask += batch['uv_mask'].numpy().sum(axis=0)
+
+    # avg_jumpsuit = sum_uv / sum_mask[np.newaxis, :, :]
+    # sum_mask[sum_mask == 0] = 1
+
+    # x = np.transpose(avg_jumpsuit, [1,2,0])
+    # x = uv2mesh(avg_jumpsuit[tf.newaxis,:,:,:])[0]
+    # go.Figure(data=[
+    #         go.Mesh3d(
+    #             # 8 vertices of a cube
+    #             x=x[:,0],
+    #             y=x[:,1],
+    #             z=x[:,2],
+    #             # i, j and k give the vertices of triangles
+    #             i = F[:,0],
+    #             j = F[:,1],
+    #             k = F[:,2]
+    #         )]).show()
+
+    # np.save('custom_data/jumpsuit_20_avg_uv.npy', np.transpose(avg_jumpsuit, [1,2,0]))
+
+    # print()
+
+
+
+
+
+    uv_res = 256
+
+    def m2uv_array(m2uv):
+        result = np.zeros((14475, 2), dtype=np.int32)
+        for k, v in m2uv.items():
+            if len(v) == 2:
+                result[k] = list(v)
+            elif len(v) == 1:
+                x = v.pop()
+                result[k] = [x, x]
+            else:
+                assert False
+
+        return result
+
+    def unique_uv_map_mask(uv_map_discreet):
+        mask = np.zeros((uv_map_discreet.shape[0]), dtype=bool)
+        temp = np.zeros((uv_res * 2, uv_res))
+        for p in uv_map_discreet:
+            temp[p[0], p[1]] += 1
+        # temp = temp == 1
+        for i, p in enumerate(uv_map_discreet):
+            if temp[p[0], p[1]] == 1:
+                mask[i] = True
+            elif temp[p[0], p[1]] == 2:
+                mask[i] = True
+                temp[p[0], p[1]] = 0
+        return mask
+
+    def uv_to_mesh(uv, m2uv):
+        res = uv.shape[0]
+        mesh = np.zeros((14475, 3))
+        uv_pixels = uv_map_discreet[m2uv]
+        # counter = np.ones((14475))
+        # counter[m2uv[:, 0] == m2uv[:, 1]] = 2
+        mesh = np.mean(uv[uv_pixels[:, :, 0], uv_pixels[:, :, 1], :], axis=1) # / counter[:, np.newaxis]
+        return mesh
+
+
+    uv_map = np.load('custom_utils/stretched_uv_map_3.npy')
+    m2uv = np.load('custom_utils/m2uv.npy', allow_pickle=True).item()
+    m2uv = m2uv_array(m2uv)
+    uv_mask = np.load('custom_utils/mask_256x512.npy')
+    curvature_label_type = '6_25'
+
+    uv_res = 256
+    uv_mask = (cv2.resize(uv_mask * 255.0, (uv_res, uv_res * 2)) / 255).astype(bool)[:, :, np.newaxis]
+    uv_mask = cv2.dilate(uv_mask.astype(np.float32), cv2.getStructuringElement(cv2.MORPH_RECT, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=2).astype(bool)[:, :, np.newaxis]
+
+    uv_map_discreet = np.floor(uv_map * uv_res).astype(np.int32)
+    uv_map_mask = unique_uv_map_mask(uv_map_discreet)
+    uv_counter = np.zeros((uv_map.shape[0]))
+    uv_counter[m2uv[:, 0]] = 1
+    uv_counter[m2uv[:, 1]] += 1
+    uv_counter = uv_counter[uv_map_mask]
+
+
+
+    for i, batch in enumerate(dataset):
+        print(i)
+        batch = static_dataset.expand_batch(batch_to_dict(batch, signature_dynamic))
+
+        uv_unposed = np.squeeze(batch['uv_unposed'].numpy())
+        mesh_unposed = np.squeeze(batch['mesh_unposed'].numpy())
+        mask = np.squeeze(batch['mesh_mask'].numpy()).astype(bool)
+
+
+        values = mesh_unposed
+        uv_values = np.zeros((uv_map.shape[0], 3))
+        uv_values[m2uv[:, 0]] = values
+        uv_values[m2uv[:, 1]] += values
+        uv_counter = np.zeros((uv_map.shape[0]))
+        uv_counter[m2uv[:, 0]] = 1
+        uv_counter[m2uv[:, 1]] += 1
+        uv_values /= uv_counter[:, np.newaxis]
+        x = uv_values
+        # verts = np.concatenate((verts, np.zeros((15068, 1))), axis=-1)
+        im = ((x - x.min()) / (x.max() - x.min())) * 255
+
+
+        x = np.transpose(uv_unposed, [1,2,0])
+        x = uv2mesh(uv_unposed[tf.newaxis,:,:,:])[0].numpy()[mask, :]
+        _, faces = _maskFaces(np.zeros((14475, 3), dtype=np.float32), F, mask)
+        go.Figure(data=[
+                go.Mesh3d(
+                    # 8 vertices of a cube
+                    x=x[:,0],
+                    y=x[:,1],
+                    z=x[:,2],
+                    # i, j and k give the vertices of triangles
+                    i = faces[:,0],
+                    j = faces[:,1],
+                    k = faces[:,2]
+                )]).show()
+
+
+        print()
+
+
+
+
+    # @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)
+        # return tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(mask[:,:,tf.newaxis] * (gt - pred) ** 2, axis=-1)), axis=-1) / tf.reduce_sum(mask, axis=-1)
+        return tf.math.sqrt(tf.reduce_sum(mask[:,:,tf.newaxis] * (gt - pred) ** 2, axis=-1))
 
     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'])
+        mesh_reconstruct = uv2mesh(batch['uv_posed'])
-        all_error.append(err.numpy())
+        err = euclid_error(batch['mesh_mask'], batch['mesh_posed'], uv2mesh(batch['uv_posed']))
+
+        i = 296
+        err_normed = (err[i].numpy() - err[i].numpy().min()) / (err[i].numpy().max() - err[i].numpy().min())
+        colors = np.array([cm.color('jet', 1 - i) for i in err_normed])
+
+        _, faces = _maskFaces(batch['mesh_posed'][i].numpy(), F, batch['mesh_mask'][i].numpy().astype(bool))
+        x = batch['mesh_posed'][i].numpy()[batch['mesh_mask'][i].numpy().astype(bool), :]
+        y = mesh_reconstruct[i].numpy()[batch['mesh_mask'][i].numpy().astype(bool), :]
+        colors = colors[batch['mesh_mask'][i].numpy().astype(bool), :]
+        fig = go.Figure(data=[
+            go.Mesh3d(
+                # 8 vertices of a cube
+                x=x[:,0],
+                y=x[:,1],
+                z=x[:,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=y[:,0],
+                y=y[:,1],
+                z=y[:,2],
+                # i, j and k give the vertices of triangles
+                i = faces[:,0],
+                j = faces[:,1],
+                k = faces[:,2],
+            )])
+        fig.show()
+
+        fig = ex.scatter_3d(x=x[:,0], y=x[:,1], z=x[:,2])
+        fig.update_traces(marker={'color': colors, 'size': 5})
+        fig.show()
+
+        all_error.append(err.numpy().sum())
         # b = nest_conditions(static_dataset.expand_batch(batch_to_dict(d, signature_dynamic)))
         print(len(all_error))
 
@@ -232,25 +442,7 @@ if __name__ == '__main__':
 
     from losses import collision_loss as collision_penalty
     from losses import collision_penalty as collision_penalty_2
-    from custom_utils.custom_utils import UV2Mesh
-    import matplotlib.pyplot as plt
 
-    def plot_3d(x, y=None):
-        fig = plt.figure(figsize=(15, 15))
-        ax = fig.add_subplot(111, projection='3d')
-        ax.set_xlim3d(-1, 1)
-        ax.set_ylim3d(-1, 1)
-        ax.set_zlim3d(-1, 1)
-        ax.set_xlabel("X")
-        ax.set_ylabel("Y")
-        ax.set_zlabel("Z")
-        ax.grid(False)
-        ax.scatter(x[:, 0], x[:, 1], x[:, 2], s=5)
-        if y is not None:
-            ax.scatter(y[:, 0], y[:, 1], y[:, 2], s=5)
-        plt.savefig('temp.jpg')
-
-    uv2mesh = UV2Mesh()
 
     for d in dataset:
         d

docker/Dockerfile

@@ -3,7 +3,7 @@ FROM tensorflow/tensorflow:2.9.0-gpu
 RUN apt-get update
 RUN apt-get upgrade -y
 RUN apt-get install git ffmpeg libsm6 libxext6 libopenexr-dev zip unzip htop nano -y
-RUN pip install numpy matplotlib opencv-python wandb OpenEXR tensorflow-graphics scipy plotly pandas tensorboardX opencv-python
+RUN pip install numpy matplotlib opencv-python wandb OpenEXR tensorflow-graphics scipy plotly pandas tensorboardX opencv-python cmapy
 
 
 WORKDIR /home/workspace/data

docker/requierements.txt

@@ -8,4 +8,5 @@ scipy
 plotly
 pandas
 tensorboardX
-opencv-python
+opencv-python
+cmapy

downscaling_test.py

@@ -0,0 +1,259 @@
+import os
+import time
+import argparse
+import wandb
+import random
+import tqdm
+import scipy as sp
+import scipy.interpolate as spi
+import numpy as np
+import tensorflow as tf
+from custom_utils.custom_utils import UV2Mesh, maskFacesBatch, _maskFaces
+import plotly.graph_objects as go
+import plotly.express as ex
+from plotly.subplots import make_subplots
+
+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 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 convnext.models.convnext_tf import ConvNeXt, ConvNeXtReverse, Block
+from stylegan2_ae.generator import Generator
+from custom_utils.graph_utils import Reconstructor_V4_Decoder as GNNDecoder
+from custom_utils.graph_utils import Reconstructor_v2 as SimpleReconstructor
+from dataset import DatasetFactory, StaticDataset
+from custom_utils.custom_utils import batch_to_dict
+from custom_utils.smpl import SMPL
+
+
+os.environ["CUDA_VISIBLE_DEVICES"] = "6"
+
+uv2mesh = UV2Mesh(res=256)
+
+
+def downscale_1(img):
+    results = [img]
+    for i in range(5):
+        results.append(tf.image.resize(img, size=np.array(img.shape[:2]) / (2 * (i + 1)), method=tf.image.ResizeMethod.LANCZOS5))
+    return results
+
+def downscale_2(img):
+    results = [img]
+    for i in range(5):
+        # print(np.array(img.shape[:2]) / (2 ** (i + 1)))
+        results.append(tf.image.resize(img, size=np.array(img.shape[:2]) / (2 ** (i + 1)), method=tf.image.ResizeMethod.NEAREST_NEIGHBOR, antialias=True, preserve_aspect_ratio=True))
+    return results
+
+def upscale_1(imgs):
+    results = []
+    for img in imgs:
+        results.append(tf.image.resize(img, size=(512, 256), method=tf.image.ResizeMethod.LANCZOS5))
+    return results
+
+
+def optimize(img, gt):
+    lr_schedule = tf.keras.optimizers.schedules.ExponentialDecay(1.0, decay_steps=40000, decay_rate=0.1, staircase=False)
+    optimizer = tf.keras.optimizers.SGD(learning_rate=lr_schedule)
+    img = tf.Variable(img)
+
+    @tf.function
+    def error():
+        pred = tf.image.resize(img, size=(512, 256), method=tf.image.ResizeMethod.LANCZOS5)
+        loss = 1 - tf.image.ssim_multiscale(pred + 1, gt + 1, max_val=2)
+        return loss
+
+    for i in range(1000):
+        optimizer.minimize(error, [img])
+
+        if i % 100 == 0:
+            pred = tf.image.resize(img, size=(512, 256), method=tf.image.ResizeMethod.LANCZOS5)
+            loss = 1 - tf.image.ssim_multiscale(pred + 1, gt + 1, max_val=2)
+            print('it: ' + str(i) + ' loss: ' + str(loss.numpy()))
+
+    return img
+
+
+
+def fill(img, mask):
+    coords = np.transpose(np.stack(np.nonzero((1 - mask))), (1,0))
+    img_idx = np.transpose(np.stack(np.nonzero((mask))), (1,0))
+    img_data = img[img_idx[:,0], img_idx[:,1], :]
+
+    interpolator =  spi.NearestNDInterpolator(x=img_idx, y=img_data)
+    img_2 = interpolator(coords)
+    res = np.zeros((512,256,3), dtype=np.float32)
+    res[coords[:,0], coords[:,1], :] = img_2
+    res[img_idx[:,0], img_idx[:,1], :] = img_data
+
+    return res
+
+
+def fill_uv_maps(folder, type):
+    os.makedirs(os.path.join(folder, type + '_filled'), exist_ok=True)
+    mask = np.load('custom_utils/mask_256x512.npy')
+
+    frames = os.listdir(os.path.join(folder, type))
+
+    for frame in tqdm.tqdm(frames):
+        id, _ = frame.split('_')
+        # mask = np.load(os.path.join(folder, 'uv_mask', id + '.npz')); mask = mask['uv_mask']
+        img = np.load(os.path.join(folder, type, frame)); img = img['uv_body']
+        filled = fill(img, mask)
+        np.savez_compressed(os.path.join(folder, type + '_filled', frame), uv_body_posed = filled)
+
+
+
+if __name__ == '__main__':
+    # fill_uv_maps(folder='/home/workspace/data/dataset/stylegan_jumpsuit/train', type='uv_body_posed')
+
+    # pass
+    # exit(0)
+
+    allow_memory_growth()
+    parser = argparse.ArgumentParser(description='')
+    parser.add_argument('--name', type=str)
+    args = vars(parser.parse_args())
+    
+    smpl = SMPL()
+
+
+    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_train = factory.get_dataset(folder='/home/workspace/local-data/stylegan_new/train',
+                                  batch_size=1,
+                                  signature=signature_dynamic)
+    static_dataset_train = StaticDataset(folder='/home/workspace/local-data/stylegan_new/train',
+                                   signature=signature_static)
+
+
+    step = 0
+    for batch in dataset_train:
+        # batch = batch_to_dict(batch, signature_dynamic)
+        batch = static_dataset_train.expand_batch(batch_to_dict(batch, signature_dynamic))
+
+        img = tf.transpose(tf.squeeze(batch['uv_posed']), [1,2,0]).numpy()
+        coords = np.transpose(np.stack(np.nonzero((1 - batch['uv_mask'][0]).numpy().astype(bool))), (1,0))
+        img_idx = np.transpose(np.stack(np.nonzero((batch['uv_mask'][0]).numpy().astype(bool))), (1,0))
+        img_data = img[img_idx[:,0], img_idx[:,1], :]
+
+        interpolator =  spi.NearestNDInterpolator(x=img_idx, y=img_data)
+        img_2 = interpolator(coords)
+        res = np.zeros((512,256,3))
+        res[coords[:,0], coords[:,1], :] = img_2
+        res[img_idx[:,0], img_idx[:,1], :] = img_data
+
+        img_2 = res
+
+        down_1 = downscale_1(img)
+        down_2 = downscale_1(img_2)
+        up_1 = upscale_1(down_1)
+        up_2 = upscale_1(down_2)
+
+        mask = batch['mesh_mask'].numpy().astype(bool)[0]
+        mesh_2 = [uv2mesh(tf.transpose(up_2[i], [2,0,1])[tf.newaxis,:,:,:])[0] for i in range(len(up_2))]
+        mesh_1 = [uv2mesh(tf.transpose(up_1[i], [2,0,1])[tf.newaxis,:,:,:])[0] for i in range(len(up_1))]
+        F = np.load("custom_utils/smpl_faces.npy")
+        _, faces = _maskFaces(np.ones((14475, 3)), F, mask)
+
+        # mesh_downscaled = mesh_2[-1].numpy()[mask, :]
+        # uv_optimized = optimize(down_2[-1], tf.transpose(batch['uv_posed'][0], [1,2,0]))
+        # uv_optimized = tf.image.resize(uv_optimized, size=(512, 256), method=tf.image.ResizeMethod.LANCZOS5)
+        # mesh_optimized = uv2mesh(tf.transpose(uv_optimized, [2,0,1])[tf.newaxis,:,:,:])[0].numpy()[mask, :]
+
+        # go.Figure(data=[
+        #     go.Mesh3d(
+        #         # 8 vertices of a cube
+        #         x=mesh_downscaled[:,0],
+        #         y=mesh_downscaled[:,1],
+        #         z=mesh_downscaled[:,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_optimized[:,0],
+        #         y=mesh_optimized[:,1],
+        #         z=mesh_optimized[:,2],
+        #         # i, j and k give the vertices of triangles
+        #         i = faces[:,0],
+        #         j = faces[:,1],
+        #         k = faces[:,2],
+        #     )
+        # ]).show()
+
+
+        # fig = make_subplots(rows=2, cols=6, start_cell="top-left")
+
+        # for i, d1 in enumerate(down_1):
+        #     fig.add_trace(ex.imshow((d1 + 1) * 128).data[0], row=1, col=i+1)
+
+        # for i, u1 in enumerate(up_1):
+        #     fig.add_trace(ex.imshow((u1 + 1) * 128).data[0], row=2, col=i+1)
+
+        # fig.show()
+
+        fig = make_subplots(rows=2, cols=6, start_cell="top-left")
+
+        for i, d2 in enumerate(down_2):
+            fig.add_trace(ex.imshow((d2 + 1) * 128).data[0], row=1, col=i+1)
+
+        for i, u2 in enumerate(up_2):
+            fig.add_trace(ex.imshow((u2 + 1) * 128).data[0], row=2, col=i+1)
+
+        fig.show()
+
+        fig = make_subplots(rows=2, cols=6, start_cell="top-left", specs = [ [{ 'type': 'xy' } for _ in range(6)], [{ 'type': 'scene' } for _ in range(6)] ])
+
+        for i, u2 in enumerate(up_2):
+            fig.add_trace(ex.imshow((u2 + 1) * 128).data[0], row=1, col=i+1)
+
+        # for i, m1 in enumerate(mesh_1):
+        #     m1 = m1.numpy()[mask, :]
+        #     fig.add_trace(go.Mesh3d(
+        #                     # 8 vertices of a cube
+        #                     x=m1[:,0],
+        #                     y=m1[:,1],
+        #                     z=m1[:,2],
+        #                     # i, j and k give the vertices of triangles
+        #                     i = faces[:,0],
+        #                     j = faces[:,1],
+        #                     k = faces[:,2]
+        #                 ),
+        #                 row=1, col=i+1)
+
+        for i, m2 in enumerate(mesh_2):
+            m2 = m2.numpy()[mask, :]
+            fig.add_trace(go.Mesh3d(
+                            # 8 vertices of a cube
+                            x=m2[:,0],
+                            y=m2[:,1],
+                            z=m2[:,2],
+                            # i, j and k give the vertices of triangles
+                            i = faces[:,0],
+                            j = faces[:,1],
+                            k = faces[:,2]
+                        ),
+                        row=2, col=i+1)
+
+        fig.show()
+
+        print()

downscaling_test_curv.py

@@ -0,0 +1,227 @@
+import os
+import time
+import argparse
+import wandb
+import random
+import tqdm
+import scipy as sp
+import cmapy as cm
+import scipy.interpolate as spi
+import numpy as np
+import tensorflow as tf
+from custom_utils.custom_utils import UV2MeshLabel, maskFacesBatch, _maskFaces
+import plotly.graph_objects as go
+import plotly.express as ex
+from plotly.subplots import make_subplots
+
+from utils import str_to_bool
+from tf_utils import check_tf_version, allow_memory_growth, split_gpu_for_testing
+# from losses import collision_penalty as collision_penalty
+from losses import collision_loss as collision_penalty
+from convnext.models.convnext_tf import ConvNeXt, ConvNeXtReverse, Block
+from stylegan2_ae.generator import Generator
+from custom_utils.graph_utils import Reconstructor_V4_Decoder as GNNDecoder
+from custom_utils.graph_utils import Reconstructor_v2 as SimpleReconstructor
+from dataset import DatasetFactory, StaticDataset
+from custom_utils.custom_utils import batch_to_dict
+from custom_utils.smpl import SMPL
+
+
+os.environ["CUDA_VISIBLE_DEVICES"] = "1"
+
+uv2mesh = UV2MeshLabel(res=256)
+
+
+def downscale_1(img):
+    results = [img]
+    for i in range(5):
+        results.append(tf.image.resize(img, size=np.array(img.shape[:2]) / (2 * (i + 1)), method=tf.image.ResizeMethod.LANCZOS5))
+    return results
+
+def downscale_2(img):
+    results = [img]
+    for i in range(5):
+        # print(np.array(img.shape[:2]) / (2 ** (i + 1)))
+        results.append(tf.image.resize(img, size=np.array(img.shape[:2]) / (2 ** (i + 1)), method=tf.image.ResizeMethod.NEAREST_NEIGHBOR, antialias=True, preserve_aspect_ratio=True))
+    return results
+
+def upscale_1(imgs):
+    results = []
+    for img in imgs:
+        results.append(tf.image.resize(img, size=(512, 256), method=tf.image.ResizeMethod.LANCZOS5))
+    return results
+
+
+def fill(img, mask):
+    coords = np.transpose(np.stack(np.nonzero((1 - mask))), (1,0))
+    img_idx = np.transpose(np.stack(np.nonzero((mask))), (1,0))
+    img_data = img[img_idx[:,0], img_idx[:,1], :]
+
+    interpolator =  spi.NearestNDInterpolator(x=img_idx, y=img_data)
+    img_2 = interpolator(coords)
+    res = np.zeros((512,256,3), dtype=np.float32)
+    res[coords[:,0], coords[:,1], :] = img_2
+    res[img_idx[:,0], img_idx[:,1], :] = img_data
+
+    return res
+
+
+def fill_uv_maps(folder, type):
+    os.makedirs(os.path.join(folder, type + '_filled'), exist_ok=True)
+    mask = np.load('custom_utils/mask_256x512.npy')
+
+    frames = os.listdir(os.path.join(folder, type))
+
+    for frame in tqdm.tqdm(frames):
+        id, _ = frame.split('_')
+        # mask = np.load(os.path.join(folder, 'uv_mask', id + '.npz')); mask = mask['uv_mask']
+        img = np.load(os.path.join(folder, type, frame)); img = img['uv_body']
+        filled = fill(img, mask)
+        np.savez_compressed(os.path.join(folder, type + '_filled', frame), uv_body_posed = filled)
+
+
+
+if __name__ == '__main__':
+    # fill_uv_maps(folder='/home/workspace/data/dataset/stylegan_jumpsuit/train', type='uv_body_posed')
+
+    # pass
+    # exit(0)
+
+    allow_memory_growth()
+    parser = argparse.ArgumentParser(description='')
+    parser.add_argument('--name', type=str)
+    args = vars(parser.parse_args())
+    
+    smpl = SMPL()
+
+
+    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,)
+                          'uv_curvature': (512, 256),
+                          'mesh_curvature':  (14475)
+                        }
+    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_train = factory.get_dataset(folder='/home/workspace/local-data/stylegan_new/train',
+                                  batch_size=1,
+                                  signature=signature_dynamic)
+    static_dataset_train = StaticDataset(folder='/home/workspace/local-data/stylegan_new/train',
+                                   signature=signature_static)
+
+
+    step = 0
+    for batch in dataset_train:
+        # batch = batch_to_dict(batch, signature_dynamic)
+        batch = static_dataset_train.expand_batch(batch_to_dict(batch, signature_dynamic))
+
+        img = tf.squeeze(batch['uv_curvature']).numpy()
+        coords = np.transpose(np.stack(np.nonzero((1 - batch['uv_mask'][0]).numpy().astype(bool))), (1,0))
+        img_idx = np.transpose(np.stack(np.nonzero((batch['uv_mask'][0]).numpy().astype(bool))), (1,0))
+        img_data = img[img_idx[:,0], img_idx[:,1]]
+
+        interpolator =  spi.NearestNDInterpolator(x=img_idx, y=img_data)
+        img_2 = interpolator(coords)
+        res = np.zeros((512,256))
+        res[coords[:,0], coords[:,1]] = img_2
+        res[img_idx[:,0], img_idx[:,1]] = img_data
+
+        img_2 = res
+
+        i = 0
+        down_curv = tf.image.resize(img_2[:,:,tf.newaxis], size=np.array(img_2.shape[:2]) / (2 * (i + 1)), method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)
+        down_curv = tf.image.resize(down_curv, size=(512, 256), method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)[:,:,0]
+
+
+
+        img_posed = tf.transpose(tf.squeeze(batch['uv_posed']), [1,2,0]).numpy()
+        coords = np.transpose(np.stack(np.nonzero((1 - batch['uv_mask'][0]).numpy().astype(bool))), (1,0))
+        img_idx = np.transpose(np.stack(np.nonzero((batch['uv_mask'][0]).numpy().astype(bool))), (1,0))
+        img_data = img_posed[img_idx[:,0], img_idx[:,1], :]
+
+        interpolator =  spi.NearestNDInterpolator(x=img_idx, y=img_data)
+        img_posed_2 = interpolator(coords)
+        res = np.zeros((512,256,3))
+        res[coords[:,0], coords[:,1], :] = img_posed_2
+        res[img_idx[:,0], img_idx[:,1], :] = img_data
+
+        img_posed_2 = res
+
+        mesh_posed = uv2mesh(tf.transpose(img_posed_2[i], [2,0,1])[tf.newaxis,:,:,:])[0]
+
+
+
+        mask = batch['mesh_mask'].numpy().astype(bool)[0]
+        mesh_1 = uv2mesh(img_2[tf.newaxis,:,:])[0].numpy()[mask]
+        mesh_2 = uv2mesh(down_curv[tf.newaxis,:,:])[0].numpy()[mask]
+        F = np.load("custom_utils/smpl_faces.npy")
+        _, faces = _maskFaces(np.ones((14475, 3)), F, mask)
+
+        fig = make_subplots(rows=1, cols=2, start_cell="top-left")
+
+        fig.add_trace(ex.imshow((img_2) / 6).data[0], row=1, col=1)
+        fig.add_trace(ex.imshow((down_curv ) / 6).data[0], row=1, col=2)
+
+        fig.show()
+
+        mesh = tf.squeeze(batch['mesh_posed']).numpy()[mask]
+        colors = np.array([cm.color('jet', 1 - i) for i in np.linspace(0, 1, 6+1)])
+        v_color_1 = np.zeros((mesh.shape[0], 3))
+        v_color_2 = np.zeros((mesh.shape[0], 3))
+        for i in range(6+1):
+            v_color_1[mesh_1 == i, :] = colors[i]
+            v_color_2[mesh_2 == i, :] = colors[i]
+
+        mesh_copy = mesh.copy()
+        mesh_copy[:,1] += 1
+        fig = go.Figure(data=[
+            go.Mesh3d(
+                # 8 vertices of a cube
+                x=mesh[:,0],
+                y=mesh[:,1],
+                z=mesh[:,2],
+                # i, j and k give the vertices of triangles
+                i = faces[:,0],
+                j = faces[:,1],
+                k = faces[:,2],
+                vertexcolor=v_color_1,
+                text='original'
+            ),
+            go.Mesh3d(
+                # 8 vertices of a cube
+                x=mesh_copy[:,0],
+                y=mesh_copy[:,1],
+                z=mesh_copy[:,2],
+                # i, j and k give the vertices of triangles
+                i = faces[:,0],
+                j = faces[:,1],
+                k = faces[:,2],
+                vertexcolor=v_color_2,
+                text='rescaled'
+            )
+        ])
+
+        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],),)
+                            )
+
+        fig.show()
+
+
+        print()

load_models.py

@@ -2,10 +2,13 @@ import tensorflow as tf
 
 
 def load_generator(g_params=None, is_g_clone=False, ckpt_dir=None, custom_cuda=True):
-    if custom_cuda:
+    # if custom_cuda:
-        from stylegan2.generator import Generator
+    #     from stylegan2.generator import Generator
-    else:
+    # else:
-        from stylegan2_ref.generator import Generator
+    #     from stylegan2_ref.generator import Generator
+    from stylegan2.generator import ConvNeXtReverseConditional as Generator
+
+
 
     if g_params is None:
         g_params = {
@@ -23,7 +26,7 @@ def load_generator(g_params=None, is_g_clone=False, ckpt_dir=None, custom_cuda=T
     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)
+    generator = Generator()
     _ = generator([test_latent, test_conditions])
 
     if ckpt_dir is not None:
@@ -44,7 +47,8 @@ def load_discriminator(d_params=None, ckpt_dir=None, custom_cuda=True):
     # else:
     #     from stylegan2_ref.discriminator import Discriminator
 
-    from custom_utils.graph_utils import GNNDiscriminator as Discriminator
+    # from custom_utils.graph_utils import GNNDiscriminator as Discriminator
+    from convnext.models.convnext_tf import ConvNeXtDiscriminator as Discriminator
 
 
     res = d_params['resolutions'][-1]
@@ -52,7 +56,8 @@ def load_discriminator(d_params=None, ckpt_dir=None, custom_cuda=True):
     test_images = tf.ones((batch_size, 3, res, int(res/2)), dtype=tf.float32)
 
     # build discriminator model
-    discriminator = Discriminator(batch_size, 32)
+    # discriminator = Discriminator(batch_size, 32)
+    discriminator = Discriminator(depths=[3, 3, 27, 3], dims=[96, 192, 256, 8], include_top=False)
     _ = discriminator(test_images)
 
     if ckpt_dir is not None:

losses.py

@@ -6,6 +6,11 @@ import scipy.io as sio
 
 from scipy.spatial import cKDTree
 from tensorflow_graphics.geometry.representation.mesh.normals import vertex_normals
+from tensorflow_graphics.geometry.convolution.graph_convolution import edge_convolution_template
+
+from custom_utils.custom_utils import maskFacesBatch, _maskFaces
+from custom_utils.graph_utils import sparse_to_tensor
+
 
 def d_logistic(real_images, generator, discriminator, z_dim, conditions):
     batch_size = tf.shape(real_images)[0]
@@ -202,6 +207,57 @@ def uv_mask_loss(pred, uv_mask):
     return invalid_count / total_count
 
 
+
+
+def merge_verts(Vs, batchDim):
+    V = tf.concat(tf.unstack(Vs, num=batchDim), axis=0)
+    return V
+    
+
+def _bend_loss_single(x, F, L, mask):
+    verts = tf.boolean_mask(x, mask, axis=0)
+    VN = vertex_normals(verts, F)
+    bend = edge_convolution_template(
+        VN, L, sizes=None,
+        edge_function=lambda x,y: (x - y)**2,
+        reduction='weighted',
+        edge_function_kwargs={}
+    )
+    loss = tf.reduce_sum(bend)
+    return loss
+
+def bend_loss(x, F, L, mask):
+    # x: predicted outfits
+    # F: faces
+    # L: laplacian
+    loss = tf.stack([_bend_loss_single(x[i], F[i], L[i], mask[i]) for i in range(x.shape[0])])
+    return loss
+
+
+# def bend_loss(x, F, L, mask):
+#     # x: predicted outfits
+#     # F: faces
+#     # L: laplacian
+#     # x = x * mask[:,:,tf.newaxis]
+#     batch_size = x.shape[0]
+#     verts, faces = maskFacesBatch(x, F, mask)
+#     nr_valid_verts = [tf.reduce_sum(mask[i]) for i in range(batch_size)]
+#     x = tf.concat(verts, axis=0)
+#     F = tf.concat([faces[i] + tf.reduce_sum(nr_valid_verts[:i]) for i in range(1, batch_size)])
+#     F = tf.concat((faces[0], F), axis=0)
+#     VN = vertex_normals(x, F)
+#     bend = edge_convolution_template(
+#         VN, L, sizes=None,
+#         edge_function=lambda x,y: (x - y)**2,
+#         reduction='weighted',
+#         edge_function_kwargs={}
+#     )
+#     loss = tf.reduce_sum(bend, axis=[1])
+#     loss = tf.reshape(loss, (batch_size, 14475))
+#     loss = tf.reduce_mean(loss, axis=1)
+#     return loss
+
+
 if __name__ == '__main__':
     V = tf.ones((4, 14475, 3))
     B = tf.zeros((4, 14475, 3))

stylegan2/generator.py

@@ -1,10 +1,19 @@
 import tensorflow as tf
+import numpy as np
 
 from stylegan2.utils import lerp
 from stylegan2.layers.mapping_block import Mapping
 from stylegan2.layers.synthesis_block import Synthesis
 
 
+from tensorflow.keras.layers import (Conv2D, Conv2DTranspose, Dense, DepthwiseConv2D,
+                                     GlobalAveragePooling2D, Layer,
+                                     LayerNormalization, Flatten, Reshape)
+# from stylegan2.layers.mapping_block import Mapping
+from stylegan2.layers.embedding import PoseShapeEmbedding, GarmentEmbedding
+from convnext.models.convnext_tf import Block, DropPath
+
+
 class Generator(tf.keras.Model):
     def __init__(self, g_params, **kwargs):
         super(Generator, self).__init__(**kwargs)
@@ -115,3 +124,86 @@ class Generator(tf.keras.Model):
 
     def compute_output_shape(self, input_shape):
         return input_shape[0][0], 3, self.resolutions[-1], self.resolutions[-1]
+
+
+
+
+class ConvNeXtReverseConditional(tf.keras.Model):
+    r""" ConvNeXt
+        A Tensorflow keras impl of : `A ConvNet for the 2020s`  -
+          https://arxiv.org/pdf/2201.03545.pdf
+
+    Args:
+        num_classes (int): Number of classes for classification head. Default: 1000
+        depths (tuple(int)): Number of blocks at each stage. Default: [3, 3, 9, 3]
+        dims (int): Feature dimension at each stage. Default: [96, 192, 384, 768]
+        include_top (bool): whether to add head or just use it as feature extractor. Default: True
+        drop_path_rate (float): Stochastic depth rate. Default: 0.
+        layer_scale_init_value (float): Init value for Layer Scale. Default: 1e-6.
+        head_init_scale (float): Init scaling value for classifier weights and biases. Default: 1.
+    """
+
+    def __init__(self, 
+                 depths=[3, 3, 9, 3], dims=[64, 64, 32, 4],
+                 drop_path_rate=0., layer_scale_init_value=1e-6, head_init_scale=1.,
+                 ):
+        super().__init__()
+
+        self.pose_shape_embedding = PoseShapeEmbedding()
+        self.garment_embedding = GarmentEmbedding((128 * 3) - (216 + 11))
+        
+        self.reshape = Reshape((16, 8, 4))
+
+        self.upsample_layers = []  # stem and 3 intermediate downsampling conv layers
+        stem = tf.keras.Sequential([
+            Conv2DTranspose(dims[0], kernel_size=4, strides=4, padding='same'),
+            LayerNormalization(epsilon=1e-6)]
+        )
+        self.upsample_layers.append(stem)
+        for i in range(3):
+            downsample_layer = tf.keras.Sequential([
+                LayerNormalization(epsilon=1e-6),
+                Conv2DTranspose(dims[i+1], kernel_size=2, strides=2, padding='same')]
+            )
+            self.upsample_layers.append(downsample_layer)
+
+        self.stages = []  # 4 feature resolution stages, each consisting of multiple residual blocks
+        dp_rates = [x for x in np.linspace(0, drop_path_rate, sum(depths))]
+        cur = 0
+        for i in range(4):
+            stage = tf.keras.Sequential(
+                [Block(dim=dims[i], drop_path=dp_rates[cur + j],
+                       layer_scale_init_value=layer_scale_init_value, prefix=f'block{i}') for j in range(depths[i])]
+            )
+            self.stages.append(stage)
+            cur += depths[i]
+
+        self.last = Dense(3, activation='tanh')
+
+    def call(self, inputs, ret_w_broadcasted=False, truncation_psi=1.0, truncation_cutoff=None,):
+        latents, conditions = inputs
+
+        pose_shape_latent = self.pose_shape_embedding([conditions['pose_shape'][:, :72], conditions['pose_shape'][:, 72:]])
+        garment_latent = self.garment_embedding(conditions['uv_static'])
+        x = tf.concat([latents, pose_shape_latent, garment_latent], axis=1)
+        x = self.reshape(x)
+
+        for i in reversed(range(4)):
+            x = self.upsample_layers[i](x)
+            x = self.stages[i](x)
+        x = self.last(x)
+
+        x = tf.transpose(x, [0,3,1,2])
+        x = x + conditions['uv_shaped_body']
+        x = x * conditions['uv_mask'][:,tf.newaxis,:,:]
+
+        return x
+
+    def set_as_moving_average_of(self, src_net):
+        return
+
+    def update_moving_average_of_w(self, w_broadcasted):
+        return
+
+    def style_mixing_regularization(self, latents1, conditions, w_broadcasted1):
+        return None

stylegan2/layers/embedding.py

@@ -91,7 +91,7 @@ class GarmentEmbedding(tf.keras.layers.Layer):
 
         # TODO: channels_first?
 
-        self.encoder = ConvNeXt(depths=[3, 3, 27, 3], dims=[96, 192, 256, self.embed_dim], include_top=False)
+        self.encoder = ConvNeXt(depths=[3, 3, 9, 3], dims=[32, 64, 128, self.embed_dim], include_top=False)
         # self.encoder(tf.keras.Input(shape=(*input_shape, 3)))
         self.pooling = tf.keras.layers.GlobalAveragePooling2D()
 

train.py

@@ -17,8 +17,8 @@ from losses import d_logistic, d_logistic_r1_reg, g_logistic_non_saturating, g_l
 from losses import collision_loss as collision_penalty
 from custom_utils.smpl import SMPL
 
-os.environ["CUDA_VISIBLE_DEVICES"] = "0"
+# os.environ["CUDA_VISIBLE_DEVICES"] = "4"
-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')
+run = wandb.init(config=tf.compat.v1.flags.FLAGS, sync_tensorboard=True, project='stylegan', name='run_37_2', 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. . reworked data loader. per garment sum of errors. per_image_standardization for garment tempalte input. convnext discriminator, convnext generator. higher starting learning rate')
 artifact = wandb.Artifact('code', type='code')
 artifact.add_file('train.py')
 artifact.add_file('losses.py')
@@ -53,7 +53,7 @@ class Trainer(object):
         self.train_res = t_params['train_res']
         self.print_step = 100
         self.save_step = 10000
-        self.image_summary_step = 1000
+        self.image_summary_step = 500
         self.unsupervised_coeff = 1
         self.gt_loss_coeff = 1
         self.reached_max_steps = False
@@ -490,8 +490,8 @@ class Trainer(object):
         test_z, test_conditions = inputs
 
         # run networks
-        fake_images_05 = self.g_clone([test_z, test_conditions], truncation_psi=0.5, training=False)
+        fake_images_05 = self.generator([test_z, test_conditions], truncation_psi=0.5, training=False)
-        fake_images_07 = self.g_clone([test_z, test_conditions], truncation_psi=0.7, training=False)
+        fake_images_07 = self.generator([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['uv_static'], test_conditions['uv_static']], axis=2)
@@ -531,12 +531,12 @@ def plot_meshes(writer, garment_unposed, garment, mesh_gt, template, body, mask,
     collision_faces = [np.concatenate((F, faces[i] + 14475), axis=0) for i in range(len(garment))]
     collision_color = [np.concatenate((np.ones_like(body[i], dtype=np.int32)*100, np.ones_like(garment[i], dtype=np.int32)*230), axis=0) for i in range(len(garment))]
 
-    for i in range(len(garment)):
+    # for i in range(len(garment)):
-        writer.add_mesh('garment_pc_pred_' + str(i+nr_offset), vertices=garment[i][np.newaxis,:,:], global_step=step)
+    #     writer.add_mesh('garment_pc_pred_' + str(i+nr_offset), vertices=garment[i][np.newaxis,:,:], global_step=step)
-        writer.add_mesh('garment_pred_' + str(i+nr_offset), vertices=garment[i][np.newaxis,:,:], faces=faces[i][np.newaxis,:,:], global_step=step)
+    #     writer.add_mesh('garment_pred_' + str(i+nr_offset), vertices=garment[i][np.newaxis,:,:], faces=faces[i][np.newaxis,:,:], global_step=step)
-        writer.add_mesh('template_' + str(i+nr_offset), vertices=template[i][np.newaxis,:,:], faces=faces[i][np.newaxis,:,:], global_step=step)
+    #     writer.add_mesh('template_' + str(i+nr_offset), vertices=template[i][np.newaxis,:,:], faces=faces[i][np.newaxis,:,:], global_step=step)
-        writer.add_mesh('body_' + str(i+nr_offset), vertices=body[i][np.newaxis,:,:], faces=F[np.newaxis,:,:], global_step=step)
+    #     writer.add_mesh('body_' + str(i+nr_offset), vertices=body[i][np.newaxis,:,:], faces=F[np.newaxis,:,:], global_step=step)
-        writer.add_mesh('collision_' + str(i+nr_offset), vertices=collision[i][np.newaxis,:,:], faces=collision_faces[i][np.newaxis,:,:], global_step=step, colors=collision_color[i][np.newaxis,:,:])
+    #     writer.add_mesh('collision_' + str(i+nr_offset), vertices=collision[i][np.newaxis,:,:], faces=collision_faces[i][np.newaxis,:,:], global_step=step, colors=collision_color[i][np.newaxis,:,:])
     
     for i in range(len(garment)):
         plotly_fig = plot_3d(mesh_gt[i], garment_unposed[i], faces[i], collision[i], collision_faces[i], collision_color[i])
@@ -576,8 +576,8 @@ def main():
     train_resolutions, train_g_featuremaps = filter_resolutions_featuremaps(resolutions, g_featuremaps, args['train_res'])
     _, train_d_featuremaps = filter_resolutions_featuremaps(resolutions, d_featuremaps, args['train_res'])
     g_params = {
-        'z_dim': 256,
+        'z_dim': 128,
-        'w_dim': 256,
+        'w_dim': 128,
         'labels_dim': -1,
         'n_mapping': 8,
         'resolutions': train_resolutions,
@@ -644,7 +644,7 @@ def main():
             'd_params': d_params,
 
             # training params
-            'g_opt': {'learning_rate': 0.002, 'beta1': 0.0, 'beta2': 0.99, 'epsilon': 1e-08, 'reg_interval': 8},
+            'g_opt': {'learning_rate': 0.001, 'beta1': 0.0, 'beta2': 0.99, 'epsilon': 1e-08, 'reg_interval': 1e10},
             'd_opt': {'learning_rate': 0.002, 'beta1': 0.0, 'beta2': 0.99, 'epsilon': 1e-08, 'reg_interval': 16},
             'batch_size': global_batch_size,
             'n_total_image': 25000000,
@@ -652,7 +652,7 @@ def main():
             'train_res': args['train_res'],
         }
 
-        trainer = Trainer(training_parameters, name=f'stylegan2-test-35-{args["train_res"]}x{args["train_res"]}')
+        trainer = Trainer(training_parameters, name=f'stylegan2-test-37_2-{args["train_res"]}x{args["train_res"]}')
         trainer.train(dist_dataset, static_dataset, signature_dynamic, strategy)
     return
 

vae.py

@@ -0,0 +1,414 @@
+import os
+import time
+import argparse
+import wandb
+import random
+import numpy as np
+import scipy as sp
+import tensorflow as tf
+from custom_utils.custom_utils import UV2Mesh, maskFacesBatch
+from tensorboardX import SummaryWriter
+import plotly.graph_objects as go
+import plotly.express as ex
+
+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 losses import d_logistic, d_logistic_r1_reg, g_logistic_non_saturating, g_logistic_ns_pathreg, edge_loss, uv_mask_loss, bend_loss
+# from losses import collision_penalty as collision_penalty
+from losses import collision_loss as collision_penalty
+from convnext.models.convnext_tf import ConvNeXt, ConvNeXtReverse, Block
+from stylegan2_ae.generator import Generator
+from custom_utils.graph_utils import Reconstructor_V4_Decoder as GNNDecoder, sparse_to_tensor
+from custom_utils.graph_utils import Reconstructor_v2 as SimpleReconstructor
+from dataset import DatasetFactory, StaticDataset
+from custom_utils.custom_utils import batch_to_dict
+from custom_utils.smpl import SMPL
+
+name='vae_11'
+os.environ["CUDA_VISIBLE_DEVICES"] = "1"
+run = wandb.init(project='vae', name=name,
+         notes='vae, ms_ssim with range 3. with l1 loss, larger latent_dim=1024. 4 dowscaling steps, 4 upscaling, the image resize for final resolution. upscale from 256*128')
+artifact = wandb.Artifact('code', type='code')
+artifact.add_file('vae.py')
+artifact.add_file('dataset.py')
+artifact.add_file('losses.py')
+artifact.add_file('convnext/models/convnext_tf.py')
+run.log_artifact(artifact)
+
+uv2mesh = UV2Mesh(res=256)
+
+def plot_3d(mesh_gt, mesh_pred_unposed, faces, collision, collision_faces, collision_color):
+    mesh_pred_unposed[:,1] += 1
+    collision[:,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_unposed[:,0],
+            y=mesh_pred_unposed[:,1],
+            z=mesh_pred_unposed[:,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=collision[:,0],
+            y=collision[:,1],
+            z=collision[:,2],
+            # i, j and k give the vertices of triangles
+            i = collision_faces[:,0],
+            j = collision_faces[:,1],
+            k = collision_faces[:,2],
+            vertexcolor=collision_color
+        )
+    ])
+
+    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 Bottleneck(tf.keras.Model):
+    
+    def __init__(self, dim):
+        super(Bottleneck, self).__init__()
+
+        self.bottleneck = tf.keras.Sequential(
+                [Block(dim=dim, drop_path=0,
+                       layer_scale_init_value=1e-6, prefix=f'bottleneck{j}') for j in range(5)]
+            )
+
+    def call(self, inputs, training=None):
+        x = inputs
+
+        x = self.bottleneck(x, training=training)
+        
+        return x
+
+
+class Encoder(tf.keras.Model):
+    
+    def __init__(self, latent_dim):
+        super(Encoder, self).__init__()
+
+        self.encoder = ConvNeXt(depths=[3, 3, 9, 3], dims=[32, 64, 128, 128], include_top=False)
+        self.flatten = tf.keras.layers.Flatten()
+        self.mean = tf.keras.layers.Dense(latent_dim)
+        self.log_var = tf.keras.layers.Dense(latent_dim)
+
+    def call(self, inputs, training=None):
+        x = inputs
+
+        x = self.encoder(x)
+        x = self.flatten(x)
+        mean = self.mean(x)
+        log_var = self.log_var(x)
+
+        return mean, log_var
+
+
+class Decoder(tf.keras.Model):
+    
+    def __init__(self):
+        super(Decoder, self).__init__()
+
+        self.dense = tf.keras.layers.Dense(8 * 4 * 512)
+        self.reshape = tf.keras.layers.Reshape((8, 4, 512))
+        self.decoder = ConvNeXtReverse(depths=list(reversed([3, 3, 9, 3])), dims=list(reversed([128, 128, 64, 16])))
+
+
+    def call(self, inputs, training=None):
+        x = inputs
+
+        x = self.dense(x)
+        x = self.reshape(x)
+        x = self.decoder(x)
+
+        return x
+
+
+class VAE(tf.keras.Model):
+    
+    def __init__(self, latent_dim=1024):
+        super(VAE, self).__init__()
+
+
+        self.encoder = Encoder(latent_dim)
+        self.decoder = Decoder()
+
+
+    @tf.function
+    def sample(self, eps=None):
+        if eps is None:
+            eps = tf.random.normal(shape=(4, self.latent_dim))
+        return self.decode(eps, apply_sigmoid=True)
+
+
+    def reparameterize(self, mean, logvar):
+        eps = tf.random.normal(shape=mean.shape)
+        return eps * tf.exp(logvar * .5) + mean
+
+
+    def decode(self, z, apply_sigmoid=False):
+        logits = self.decoder(z)
+        if apply_sigmoid:
+            probs = tf.sigmoid(logits)
+            return probs
+        return logits
+
+
+    def call(self, inputs, training=None):
+        x, uv_mask = inputs
+
+        x = tf.transpose(x, [0,2,3,1])
+
+        mean, logvar = self.encoder(x)
+        z = self.reparameterize(mean, logvar)
+        x = self.decoder(z)
+
+
+        # x = tf.transpose(x, [0,3,1,2])
+        # x = x * uv_mask[:,tf.newaxis,:,:]
+        x = tf.image.resize(x, (512, 256), tf.image.ResizeMethod.LANCZOS5)
+        x = x * uv_mask[:,:,:,tf.newaxis]
+
+        return x, mean, logvar, z
+
+
+def log_normal_pdf(sample, mean, logvar, raxis=1):
+    log2pi = tf.math.log(2. * np.pi)
+    res = tf.reduce_sum(-.5 * ((sample - mean) ** 2. * tf.exp(-logvar) + logvar + log2pi), axis=raxis)
+    return res
+
+
+@tf.function
+def train_step(batch):
+
+    with tf.GradientTape() as tape:
+
+        pred, mean, logvar, z = model([batch['uv_unposed'], batch['uv_mask']])
+
+        mesh_pred = uv2mesh(tf.transpose(pred, [0,3,1,2]))
+        mesh_orig = batch['mesh_unposed']
+
+
+        e_loss = 1e-3 * tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']))
+        c_loss = 1e-3 * tf.reduce_mean(collision_penalty(mesh_pred, batch['mesh_body_posed'], batch['mesh_mask']))
+        l1_loss = tf.reduce_sum(tf.reduce_sum((tf.math.abs(tf.transpose(batch['uv_unposed'], [0,2,3,1]) - pred)), axis=-1) * batch['uv_mask'], axis=[1, 2])
+        # ms_ssim = tf.image.ssim_multiscale(tf.transpose(pred, [0,2,3,1]) + 1, tf.transpose(batch['uv_unposed'], [0,2,3,1]) + 1, max_val=2)
+        ms_ssim = tf.image.ssim_multiscale(pred + 1, tf.transpose(batch['uv_unposed'], [0,2,3,1]) + 1, max_val=3)
+        # ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale((tf.transpose(pred, [0,2,3,1]) + 1) * batch['uv_mask'][:,:,:,tf.newaxis], (tf.transpose(batch['uv_unposed'], [0,2,3,1]) + 1) * batch['uv_mask'][:,:,:,tf.newaxis], max_val=2))
+        rfft_loss = 0 #tf.reduce_mean(tf.reduce_sum(tf.math.abs(tf.signal.rfft3d(pred) - tf.signal.rfft3d(batch['uv_unposed'])), axis=[1,2,3]))
+        fft_loss = tf.reduce_sum(tf.math.abs(tf.signal.fft3d(tf.complex(pred, tf.zeros_like(pred))) - tf.signal.fft3d(tf.complex(tf.transpose(batch['uv_unposed'], [0,2,3,1]), tf.zeros_like(tf.transpose(batch['uv_unposed'], [0,2,3,1]))))), axis=[1,2,3])
+        fft_loss_pc, fft_loss_pc_low, fft_loss_pc_high = 0,0,0 #fft_loss_per_channel(pred, batch['uv_unposed'])
+        tv_loss = tf.image.total_variation(pred)
+
+        reconstr_loss = l1_loss
+        logpx_z = -reconstr_loss
+        logpz = log_normal_pdf(z, 0., 0.)
+        logqz_x = log_normal_pdf(z, mean, logvar)
+        loss = -tf.reduce_mean(logpx_z + logpz - logqz_x)
+        ms_ssim = tf.reduce_mean(ms_ssim)
+        l1_loss = tf.reduce_mean(l1_loss)
+        fft_loss = tf.reduce_mean(fft_loss)
+        reconstr_loss = tf.reduce_mean(reconstr_loss)
+        tv_loss = tf.reduce_mean(tv_loss)
+
+
+    mesh_orig_2 = uv2mesh(batch['uv_unposed'])
+    metric = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
+    metric_2 = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig_2 - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
+    euclid_posed = 0 # tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_skinned - batch['mesh_posed']) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
+    gradients = tape.gradient(loss, model.trainable_variables)
+    # gradients = [tf.clip_by_norm(g, 1.0) for g in gradients]
+    # gradients = [tf.clip_by_value(g, -1.0, 1.0) for g in gradients]
+    optimizer.apply_gradients(zip(gradients, model.trainable_variables))
+    return { 'train/loss': loss, 'train/edge_loss': e_loss, 'train/collision_loss': c_loss, 'train/l1_loss': l1_loss, 'train/euclid_error': metric, 'train/ms_ssim': ms_ssim,
+             'train/fft_loss': fft_loss, 'train/euclid_error_reconst': metric_2, 'train/rfft_loss': rfft_loss, 'train/fft_oc_low': fft_loss_pc_low, 'train/fft_pc_high': fft_loss_pc_high,
+             'train/fft_loss_per_channel': fft_loss_pc, 'train/euclid_error_posed': euclid_posed, 'train/reconstr_loss': reconstr_loss, 'train/tv_loss': tv_loss }
+        
+
+
+@tf.function
+def valid_step(batch):
+    pred, mean, logvar, z = model([batch['uv_unposed'], batch['uv_mask']])
+
+    mesh_pred = uv2mesh(tf.transpose(pred, [0,3,1,2]))
+    mesh_orig = batch['mesh_unposed']
+    e_loss = 1e-2 * tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']))
+    c_loss = 1e-2 * tf.reduce_mean(collision_penalty(mesh_pred, batch['mesh_body_posed'], batch['mesh_mask']))
+    l1_loss = tf.reduce_sum(tf.reduce_sum((tf.math.abs(tf.transpose(batch['uv_unposed'], [0,2,3,1]) - pred)), axis=-1) * batch['uv_mask'], axis=[1, 2])
+    b_loss = 0 #1e-3 * bend_loss(mesh_pred, smpl_faces, laplacian_loss, batch['mesh_mask'])
+    # ms_ssim = tf.image.ssim_multiscale(tf.transpose(pred, [0,2,3,1]) + 1, tf.transpose(batch['uv_unposed'], [0,2,3,1]) + 1, max_val=2)
+    ms_ssim = tf.image.ssim_multiscale(pred + 1, tf.transpose(batch['uv_unposed'], [0,2,3,1]) + 1, max_val=3)
+    fft_loss = 0 #tf.reduce_sum(tf.math.abs(tf.signal.fft3d(tf.complex(pred, tf.zeros_like(pred))) - tf.signal.fft3d(tf.complex(tf.transpose(batch['uv_unposed'], [0,2,3,1]), tf.zeros_like(tf.transpose(batch['uv_unposed'], [0,2,3,1]))))), axis=[1,2,3])
+    # mesh_orig_2 = uv2mesh(batch['uv_unposed'])
+    metric = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
+    metric_2 = 0 #tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig_2 - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
+  
+    reconstr_loss = l1_loss
+    logpx_z = -reconstr_loss
+    logpz = log_normal_pdf(z, 0., 0.)
+    logqz_x = log_normal_pdf(z, mean, logvar)
+    loss = -tf.reduce_mean(logpx_z + logpz - logqz_x)
+    ms_ssim = tf.reduce_mean(ms_ssim)
+    l1_loss = tf.reduce_mean(l1_loss)
+    fft_loss = tf.reduce_mean(fft_loss)
+    b_loss = tf.reduce_mean(b_loss)
+    reconstr_loss = tf.reduce_mean(reconstr_loss)
+
+    return { 'valid/loss': loss, 'valid/edge_loss': e_loss, 'valid/collision_loss': c_loss, 'valid/l1_loss': l1_loss, 'valid/euclid_error': metric, 'valid/ms_ssim': ms_ssim,
+             'valid/fft_loss': fft_loss, 'valid/euclid_error_reconst': metric_2, 'train/bend_loss': b_loss, 'train/reconstr_loss': reconstr_loss
+            }
+
+
+def plot_everything(batch: dict, type: str):
+
+    pred, mean, logvar, z = model([batch['uv_unposed'], batch['uv_mask']])
+    
+    mesh_pred = uv2mesh(tf.transpose(pred, [0,3,1,2]))
+    mesh_orig = batch['mesh_unposed']
+    
+    orig_img = tf.transpose(batch['uv_unposed'], [0,2,3,1])
+    pred_img = pred
+    img = tf.concat((orig_img, pred_img), axis=2)
+    imgs = [wandb.Image((img[i] + 1) / 2.0) for i in range(4)]
+
+    F = np.load("custom_utils/smpl_faces.npy")
+    mask = batch['mesh_mask'].numpy().astype(bool)
+    garment_unposed, faces = maskFacesBatch(mesh_pred.numpy(), F, mask)
+    gt = [mesh_orig[i].numpy()[mask[i], :] for i in range(pred.shape[0])]
+    
+    body = batch['mesh_body_posed'].numpy()
+    mesh_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred).numpy()
+    garment_posed = [mesh_posed[i, mask[i], :] for i in range(body.shape[0])]
+    collision = [np.concatenate((body[i], garment_posed[i]), axis=0) for i in range(len(garment_posed))]
+    collision_faces = [np.concatenate((F, faces[i] + 14475), axis=0) for i in range(len(garment_posed))]
+    collision_color = [np.concatenate((np.ones_like(body[i], dtype=np.int32)*100, np.ones_like(garment_posed[i], dtype=np.int32)*230), axis=0) for i in range(len(garment_posed))]
+
+    figs = { type + "/garment_" + str(i): plot_3d(gt[i], garment_unposed[i], faces[i], collision[i], collision_faces[i], collision_color[i]) for i in range(4) }
+    
+    return { **figs, **{ type + '/images': imgs } }
+
+
+
+
+
+def repeat_faces(faces, times):
+    idx_offset = np.linspace(0, 14475  * (times-1), times).astype(np.int32)
+    idx_offset = np.repeat(idx_offset, faces.shape[0], axis=0)
+    merged = np.tile(faces, [times,1])
+    merged += idx_offset[:, np.newaxis]
+    return merged
+
+
+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)
+
+
+if __name__ == '__main__':
+    allow_memory_growth()
+    parser = argparse.ArgumentParser(description='')
+    parser.add_argument('--name', type=str)
+    args = vars(parser.parse_args())
+
+    batch_size = 8
+
+    smpl_faces = np.load('custom_utils/smpl_faces.npy')
+    faces_loss_single = smpl_faces
+    faces_loss = repeat_faces(smpl_faces, batch_size)
+    laplacian_loss = sp.sparse.load_npz("custom_utils/laplacian/laplacian0.npz").todense().astype(np.float32)
+    # edges = np.stack(np.nonzero(laplacian_loss.todense())).transpose()
+    # laplacian_loss = sparse_to_tensor(laplacian_loss)
+    # laplacian_loss = sparse_to_tensor(repeat_laplacian(laplacian_loss, batch_size))
+    
+    smpl = SMPL()
+
+    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_train = factory.get_dataset(folder='/home/workspace/local-data/jumpsuit_20/train',
+                                  batch_size=batch_size,
+                                  signature=signature_dynamic)
+    static_dataset_train = StaticDataset(folder='/home/workspace/local-data/jumpsuit_20/train',
+                                   signature=signature_static)
+    dataset_valid = factory.get_dataset(folder='/home/workspace/local-data/jumpsuit_20/valid',
+                                  batch_size=batch_size,
+                                  signature=signature_dynamic)
+    static_dataset_valid = StaticDataset(folder='/home/workspace/local-data/jumpsuit_20/valid',
+                                   signature=signature_static)
+
+
+    model = VAE()
+    lr_schedule = tf.keras.optimizers.schedules.ExponentialDecay(0.0001, decay_steps=30000, decay_rate=0.1, staircase=False)
+    optimizer = tf.optimizers.Adam(learning_rate=lr_schedule)
+    folder = 'models/ae_logs/'
+    valid_iterator = iter(dataset_valid)
+
+    step = 0
+    for batch in dataset_train:
+        batch = static_dataset_train.expand_batch(batch_to_dict(batch, signature_dynamic))
+
+        # _, masked_faces = maskFacesBatch(batch['mesh_unposed'].numpy(), smpl_faces, batch['mesh_mask'].numpy().astype(bool))
+        # # batch['mesh_masked_verts'] = masked_verts
+        # batch['mesh_masked_faces'] = [masked_faces[i].astype(np.int32) for i in range(batch_size)]
+        # batch['mesh_mask_bool'] = batch['mesh_mask'].numpy().astype(bool)
+        # batch['laplacian'] = [sparse_to_tensor(sp.sparse.coo_matrix(laplacian_loss[batch['mesh_mask_bool'][i], :][:, batch['mesh_mask_bool'][i]])) for i in range(batch_size)]
+        # batch['laplacian'] = [tf.sparse.to_dense(sparse_to_tensor(sp.sparse.coo_matrix(laplacian_loss[batch['mesh_mask_bool'][i], :][:, batch['mesh_mask_bool'][i]]))) for i in range(batch_size)]
+        # batch['laplacian'] = [sp.sparse.coo_matrix(laplacian_loss[batch['mesh_mask_bool'][i], :][:, batch['mesh_mask_bool'][i]]) for i in range(batch_size)]
+
+        train_loss = train_step(batch)
+
+        data_valid = valid_iterator.get_next()
+        data_valid = static_dataset_valid.expand_batch(batch_to_dict(data_valid, signature_dynamic))
+        valid_loss = valid_step(data_valid)
+
+        if step % 500 == 0:
+            train_figs = plot_everything(batch, 'train')
+            valid_figs = plot_everything(data_valid, 'valid')
+            wandb.log({ **train_loss, **valid_loss, **{ 'current_step': step}, **train_figs, **valid_figs }, step=step)
+        else:
+            wandb.log({ **train_loss, **valid_loss, **{ 'current_step': step} }, step=step)
+            
+
+        print(step)
+
+        step += 1

verify_scales.py

@@ -0,0 +1,81 @@
+import os
+import sys
+import time
+import cv2
+import random
+import numpy as np
+import tensorflow as tf
+
+from dataset import DatasetFactory, StaticDataset
+from custom_utils.custom_utils import batch_to_dict, UV2Mesh, _maskFaces
+import plotly.graph_objects as go
+import plotly.express as ex
+import cmapy as cm
+
+if __name__ == '__main__':
+    os.environ["CUDA_VISIBLE_DEVICES"] = "6"
+
+    uv2mesh = UV2Mesh()
+
+    F = np.load("custom_utils/smpl_faces.npy")
+
+    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,),
+                        #   'uv_curvature': (512, 256),
+                        #   'mesh_curvature':  (14475)
+                        }
+    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/local-data/jumpsuit_20/train',
+                                  batch_size=1,
+                                  signature=signature_dynamic)
+    static_dataset = StaticDataset(folder='/home/workspace/local-data/jumpsuit_20/train',
+                                   signature=signature_static)
+
+
+    for i, batch in enumerate(dataset):
+        print(i)
+        batch = static_dataset.expand_batch(batch_to_dict(batch, signature_dynamic))
+
+        uv_unposed = np.squeeze(batch['uv_unposed'].numpy())
+        mesh_unposed = np.squeeze(batch['mesh_unposed'].numpy())
+        mask = np.squeeze(batch['mesh_mask'].numpy()).astype(bool)
+        _, faces = _maskFaces(np.ones((14475, 3)), F, mask)
+
+        mesh_unposed[:,1] += (1 * i)
+        fig = go.Figure(data=[
+            go.Mesh3d(
+                # 8 vertices of a cube
+                x=mesh_unposed[mask,0],
+                y=mesh_unposed[mask,2],
+                z=mesh_unposed[mask,1],
+                # 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=10, z=1))
+        fig.update_layout(
+            scene = dict(
+                xaxis = dict(nticks=4, range=[-1,1],),
+                            yaxis = dict(nticks=4, range=[-1,19],),
+                            zaxis = dict(nticks=4, range=[-1,1],),)
+                            )
+
+        fig.show()
+
+        print()

verify_uv_mapping.py

@@ -0,0 +1,145 @@
+import os
+import sys
+import time
+from unicodedata import name
+import cv2
+import random
+import numpy as np
+import tensorflow as tf
+
+from dataset import DatasetFactory, StaticDataset
+from custom_utils.custom_utils import batch_to_dict, UV2Mesh, _maskFaces
+import plotly.graph_objects as go
+import plotly.express as ex
+import cmapy as cm
+
+if __name__ == '__main__':
+    os.environ["CUDA_VISIBLE_DEVICES"] = "6"
+
+    uv2mesh = UV2Mesh()
+
+    F = np.load("custom_utils/smpl_faces.npy")
+
+    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,),
+                        #   'uv_curvature': (512, 256),
+                        #   'mesh_curvature':  (14475)
+                        }
+    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/local-data/jumpsuit_20/train',
+                                  batch_size=1,
+                                  signature=signature_dynamic)
+    static_dataset = StaticDataset(folder='/home/workspace/local-data/jumpsuit_20/train',
+                                   signature=signature_static)
+
+
+    uv_res = 256
+
+    def m2uv_array(m2uv):
+        result = np.zeros((14475, 2), dtype=np.int32)
+        for k, v in m2uv.items():
+            if len(v) == 2:
+                result[k] = list(v)
+            elif len(v) == 1:
+                x = v.pop()
+                result[k] = [x, x]
+            else:
+                assert False
+
+        return result
+
+    def unique_uv_map_mask(uv_map_discreet):
+        mask = np.zeros((uv_map_discreet.shape[0]), dtype=bool)
+        temp = np.zeros((uv_res * 2, uv_res))
+        for p in uv_map_discreet:
+            temp[p[0], p[1]] += 1
+        # temp = temp == 1
+        for i, p in enumerate(uv_map_discreet):
+            if temp[p[0], p[1]] == 1:
+                mask[i] = True
+            elif temp[p[0], p[1]] == 2:
+                mask[i] = True
+                temp[p[0], p[1]] = 0
+        return mask
+
+    def uv_to_mesh(uv, m2uv):
+        res = uv.shape[0]
+        mesh = np.zeros((14475, 3))
+        uv_pixels = uv_map_discreet[m2uv]
+        # counter = np.ones((14475))
+        # counter[m2uv[:, 0] == m2uv[:, 1]] = 2
+        mesh = np.mean(uv[uv_pixels[:, :, 0], uv_pixels[:, :, 1], :], axis=1) # / counter[:, np.newaxis]
+        return mesh
+
+
+    uv_map = np.load('custom_utils/stretched_uv_map_3.npy')
+    m2uv = np.load('custom_utils/m2uv.npy', allow_pickle=True).item()
+    m2uv = m2uv_array(m2uv)
+    uv_mask = np.load('custom_utils/mask_256x512.npy')
+    curvature_label_type = '6_25'
+
+    uv_res = 256
+    uv_mask = (cv2.resize(uv_mask * 255.0, (uv_res, uv_res * 2)) / 255).astype(bool)[:, :, np.newaxis]
+    uv_mask = cv2.dilate(uv_mask.astype(np.float32), cv2.getStructuringElement(cv2.MORPH_RECT, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=2).astype(bool)[:, :, np.newaxis]
+
+    uv_map_discreet = np.floor(uv_map * uv_res).astype(np.int32)
+    uv_map_mask = unique_uv_map_mask(uv_map_discreet)
+    uv_counter = np.zeros((uv_map.shape[0]))
+    uv_counter[m2uv[:, 0]] = 1
+    uv_counter[m2uv[:, 1]] += 1
+    uv_counter = uv_counter[uv_map_mask]
+
+
+
+    for i, batch in enumerate(dataset):
+        print(i)
+        batch = static_dataset.expand_batch(batch_to_dict(batch, signature_dynamic))
+
+        uv_unposed = np.squeeze(batch['uv_unposed'].numpy())
+        mesh_unposed = np.squeeze(batch['mesh_unposed'].numpy())
+        mask = np.squeeze(batch['mesh_mask'].numpy()).astype(bool)
+
+
+        values = mesh_unposed
+        uv_values = np.zeros((uv_map.shape[0], 3))
+        uv_values[m2uv[:, 0]] = values
+        uv_values[m2uv[:, 1]] += values
+        uv_counter = np.zeros((uv_map.shape[0]))
+        uv_counter[m2uv[:, 0]] = 1
+        uv_counter[m2uv[:, 1]] += 1
+        uv_values /= uv_counter[:, np.newaxis]
+        x = uv_values
+        # verts = np.concatenate((verts, np.zeros((15068, 1))), axis=-1)
+        im = ((x - x.min()) / (x.max() - x.min())) * 255
+
+
+        x = np.transpose(uv_unposed, [1,2,0])
+        x = uv2mesh(uv_unposed[tf.newaxis,:,:,:])[0].numpy()[mask, :]
+        _, faces = _maskFaces(np.zeros((14475, 3), dtype=np.float32), F, mask)
+        go.Figure(data=[
+                go.Mesh3d(
+                    # 8 vertices of a cube
+                    x=x[:,0],
+                    y=x[:,1],
+                    z=x[:,2],
+                    # i, j and k give the vertices of triangles
+                    i = faces[:,0],
+                    j = faces[:,1],
+                    k = faces[:,2]
+                )]).show()
+
+
+        print()