remove unused parts

configs/iccvw/ae_body_1_10p.py

@@ -0,0 +1,78 @@
+from utils.convnext_utils import CNBlockConfig
+
+config = dict()
+
+config['gpu'] = '5'
+config['name'] = 'body_iccvw_1_10p'
+config['notes'] = ' ' 
+config['tags'] = ['torch', 'body'] 
+config['script'] = 'scripts/ae_body_train.py'
+
+config['pretrained_model'] = ''
+
+config['batch_size'] = 4
+config['train_dir'] = '/data-local/127/cloth3d_processed'
+config['valid_dir'] = '/data-local/127/cloth3d_processed'
+config['test_dir'] = '/data-net/dataset/test_ready'
+
+config['resolution'] = 512
+config['is_skirt'] = False
+config['dataset_in_memory'] = False
+config['is_cape'] = False
+config['remove_global_orientation'] = True
+config['device'] = 'cuda'
+
+if config['is_skirt']:
+  config['train_frames'] = 'frames_train_skirt_10p.txt'
+  config['valid_frames'] = 'frames_valid_skirt.txt'
+else:
+  config['train_frames'] = 'frames_train_10p.txt'
+  config['valid_frames'] = 'frames_valid.txt'
+
+config['smpl_res'] = 14475
+if config['is_cape']:
+  config['smpl_res'] = 3638
+
+config['latent_dim'] = 128
+
+config['bottleneck_scale'] = 2
+
+config['cond_ae'] = dict()
+config['cond_ae']['device'] = config['device']
+config['cond_ae']['is_label'] = False
+config['cond_ae']['latent_dim'] = config['latent_dim']
+config['cond_ae']['kernel_size'] = 7
+config['cond_ae']['bottleneck_size'] = 16 * config['bottleneck_scale']
+
+block_setting = [
+    CNBlockConfig(16, 32, 3),
+    CNBlockConfig(32, 64, 3),
+    CNBlockConfig(64, 128, 9),
+    CNBlockConfig(128, config['cond_ae']['bottleneck_size'], 9),
+    CNBlockConfig(config['cond_ae']['bottleneck_size'], None, 3),
+]
+
+block_setting_reverse = [
+    CNBlockConfig(config['cond_ae']['bottleneck_size'], 128, 3),
+    CNBlockConfig(128, 64, 9),
+    CNBlockConfig(64, 32, 9),
+    CNBlockConfig(32, 16, 3),
+    CNBlockConfig(16, 3, 3),
+    # CNBlockConfig(3, 3, 3),
+]
+
+config['cond_ae']['encoder_block_settings'] = block_setting
+config['cond_ae']['decoder_block_settings'] = block_setting_reverse
+config['cond_ae']['activation'] = 'tanh'
+
+
+config['lr'] = 0.0001
+config['cond_ae']['upscale_output'] = False
+
+config['signature'] = { 'id',
+                        'frame_nr',
+                        'pose_shape',
+                        'uv_body_posed',
+                        'mesh_mask',
+                        'mesh_body_posed',
+                        }

configs/iccvw/ae_body_1_25p.py

@@ -0,0 +1,78 @@
+from utils.convnext_utils import CNBlockConfig
+
+config = dict()
+
+config['gpu'] = '5'
+config['name'] = 'body_iccvw_1_25p'
+config['notes'] = ' ' 
+config['tags'] = ['torch', 'body'] 
+config['script'] = 'scripts/ae_body_train.py'
+
+config['pretrained_model'] = ''
+
+config['batch_size'] = 4
+config['train_dir'] = '/data-local/127/cloth3d_processed'
+config['valid_dir'] = '/data-local/127/cloth3d_processed'
+config['test_dir'] = '/data-net/dataset/test_ready'
+
+config['resolution'] = 512
+config['is_skirt'] = False
+config['dataset_in_memory'] = False
+config['is_cape'] = False
+config['remove_global_orientation'] = True
+config['device'] = 'cuda'
+
+if config['is_skirt']:
+  config['train_frames'] = 'frames_train_skirt_25p.txt'
+  config['valid_frames'] = 'frames_valid_skirt.txt'
+else:
+  config['train_frames'] = 'frames_train_25p.txt'
+  config['valid_frames'] = 'frames_valid.txt'
+
+config['smpl_res'] = 14475
+if config['is_cape']:
+  config['smpl_res'] = 3638
+
+config['latent_dim'] = 128
+
+config['bottleneck_scale'] = 2
+
+config['cond_ae'] = dict()
+config['cond_ae']['device'] = config['device']
+config['cond_ae']['is_label'] = False
+config['cond_ae']['latent_dim'] = config['latent_dim']
+config['cond_ae']['kernel_size'] = 7
+config['cond_ae']['bottleneck_size'] = 16 * config['bottleneck_scale']
+
+block_setting = [
+    CNBlockConfig(16, 32, 3),
+    CNBlockConfig(32, 64, 3),
+    CNBlockConfig(64, 128, 9),
+    CNBlockConfig(128, config['cond_ae']['bottleneck_size'], 9),
+    CNBlockConfig(config['cond_ae']['bottleneck_size'], None, 3),
+]
+
+block_setting_reverse = [
+    CNBlockConfig(config['cond_ae']['bottleneck_size'], 128, 3),
+    CNBlockConfig(128, 64, 9),
+    CNBlockConfig(64, 32, 9),
+    CNBlockConfig(32, 16, 3),
+    CNBlockConfig(16, 3, 3),
+    # CNBlockConfig(3, 3, 3),
+]
+
+config['cond_ae']['encoder_block_settings'] = block_setting
+config['cond_ae']['decoder_block_settings'] = block_setting_reverse
+config['cond_ae']['activation'] = 'tanh'
+
+
+config['lr'] = 0.0001
+config['cond_ae']['upscale_output'] = False
+
+config['signature'] = { 'id',
+                        'frame_nr',
+                        'pose_shape',
+                        'uv_body_posed',
+                        'mesh_mask',
+                        'mesh_body_posed',
+                        }

configs/iccvw/ae_body_1_50p.py

@@ -0,0 +1,78 @@
+from utils.convnext_utils import CNBlockConfig
+
+config = dict()
+
+config['gpu'] = '5'
+config['name'] = 'body_iccvw_1_50p'
+config['notes'] = ' ' 
+config['tags'] = ['torch', 'body'] 
+config['script'] = 'scripts/ae_body_train.py'
+
+config['pretrained_model'] = ''
+
+config['batch_size'] = 4
+config['train_dir'] = '/data-local/127/cloth3d_processed'
+config['valid_dir'] = '/data-local/127/cloth3d_processed'
+config['test_dir'] = '/data-net/dataset/test_ready'
+
+config['resolution'] = 512
+config['is_skirt'] = False
+config['dataset_in_memory'] = False
+config['is_cape'] = False
+config['remove_global_orientation'] = True
+config['device'] = 'cuda'
+
+if config['is_skirt']:
+  config['train_frames'] = 'frames_train_skirt_50p.txt'
+  config['valid_frames'] = 'frames_valid_skirt.txt'
+else:
+  config['train_frames'] = 'frames_train_50p.txt'
+  config['valid_frames'] = 'frames_valid.txt'
+
+config['smpl_res'] = 14475
+if config['is_cape']:
+  config['smpl_res'] = 3638
+
+config['latent_dim'] = 128
+
+config['bottleneck_scale'] = 2
+
+config['cond_ae'] = dict()
+config['cond_ae']['device'] = config['device']
+config['cond_ae']['is_label'] = False
+config['cond_ae']['latent_dim'] = config['latent_dim']
+config['cond_ae']['kernel_size'] = 7
+config['cond_ae']['bottleneck_size'] = 16 * config['bottleneck_scale']
+
+block_setting = [
+    CNBlockConfig(16, 32, 3),
+    CNBlockConfig(32, 64, 3),
+    CNBlockConfig(64, 128, 9),
+    CNBlockConfig(128, config['cond_ae']['bottleneck_size'], 9),
+    CNBlockConfig(config['cond_ae']['bottleneck_size'], None, 3),
+]
+
+block_setting_reverse = [
+    CNBlockConfig(config['cond_ae']['bottleneck_size'], 128, 3),
+    CNBlockConfig(128, 64, 9),
+    CNBlockConfig(64, 32, 9),
+    CNBlockConfig(32, 16, 3),
+    CNBlockConfig(16, 3, 3),
+    # CNBlockConfig(3, 3, 3),
+]
+
+config['cond_ae']['encoder_block_settings'] = block_setting
+config['cond_ae']['decoder_block_settings'] = block_setting_reverse
+config['cond_ae']['activation'] = 'tanh'
+
+
+config['lr'] = 0.0001
+config['cond_ae']['upscale_output'] = False
+
+config['signature'] = { 'id',
+                        'frame_nr',
+                        'pose_shape',
+                        'uv_body_posed',
+                        'mesh_mask',
+                        'mesh_body_posed',
+                        }

configs/iccvw/ae_body_1_skirt.py

@@ -0,0 +1,78 @@
+from utils.convnext_utils import CNBlockConfig
+
+config = dict()
+
+config['gpu'] = '1'
+config['name'] = 'body_iccvw_1_skirt'
+config['notes'] = ' ' 
+config['tags'] = ['torch', 'body'] 
+config['script'] = 'scripts/ae_body_train.py'
+
+config['pretrained_model'] = ''
+
+config['batch_size'] = 4
+config['train_dir'] = '/data-local/127/cloth3d_processed'
+config['valid_dir'] = '/data-local/127/cloth3d_processed'
+config['test_dir'] = '/data-net/dataset/test_ready'
+
+config['resolution'] = 512
+config['is_skirt'] = True
+config['dataset_in_memory'] = False
+config['is_cape'] = False
+config['remove_global_orientation'] = True
+config['device'] = 'cuda'
+
+if config['is_skirt']:
+  config['train_frames'] = 'frames_train_skirt.txt'
+  config['valid_frames'] = 'frames_valid_skirt.txt'
+else:
+  config['train_frames'] = 'frames_train.txt'
+  config['valid_frames'] = 'frames_valid.txt'
+
+config['smpl_res'] = 14475
+if config['is_cape']:
+  config['smpl_res'] = 3638
+
+config['latent_dim'] = 128
+
+config['bottleneck_scale'] = 2
+
+config['cond_ae'] = dict()
+config['cond_ae']['device'] = config['device']
+config['cond_ae']['is_label'] = False
+config['cond_ae']['latent_dim'] = config['latent_dim']
+config['cond_ae']['kernel_size'] = 7
+config['cond_ae']['bottleneck_size'] = 16 * config['bottleneck_scale']
+
+block_setting = [
+    CNBlockConfig(16, 32, 3),
+    CNBlockConfig(32, 64, 3),
+    CNBlockConfig(64, 128, 9),
+    CNBlockConfig(128, config['cond_ae']['bottleneck_size'], 9),
+    CNBlockConfig(config['cond_ae']['bottleneck_size'], None, 3),
+]
+
+block_setting_reverse = [
+    CNBlockConfig(config['cond_ae']['bottleneck_size'], 128, 3),
+    CNBlockConfig(128, 64, 9),
+    CNBlockConfig(64, 32, 9),
+    CNBlockConfig(32, 16, 3),
+    CNBlockConfig(16, 3, 3),
+    # CNBlockConfig(3, 3, 3),
+]
+
+config['cond_ae']['encoder_block_settings'] = block_setting
+config['cond_ae']['decoder_block_settings'] = block_setting_reverse
+config['cond_ae']['activation'] = 'tanh'
+
+
+config['lr'] = 0.0001
+config['cond_ae']['upscale_output'] = False
+
+config['signature'] = { 'id',
+                        'frame_nr',
+                        'pose_shape',
+                        'uv_body_posed',
+                        'mesh_mask',
+                        'mesh_body_posed',
+                        }

configs/iccvw/ae_norm_1_skirt.py

@@ -0,0 +1,78 @@
+from utils.convnext_utils import CNBlockConfig
+
+config = dict()
+
+config['gpu'] = '4'
+config['name'] = 'norm_iccvw_1_skirt'
+config['notes'] = ' ' 
+config['tags'] = ['torch', 'norm'] 
+config['script'] = 'scripts/ae_norm_train.py'
+
+config['pretrained_model'] = ''
+
+config['batch_size'] = 4
+config['train_dir'] = '/data-local/127/cloth3d_processed'
+config['valid_dir'] = '/data-local/127/cloth3d_processed'
+config['test_dir'] = '/data-net/dataset/test_ready'
+
+config['resolution'] = 512
+config['is_skirt'] = True
+config['dataset_in_memory'] = False
+config['is_cape'] = False
+config['remove_global_orientation'] = True
+config['device'] = 'cuda'
+
+if config['is_skirt']:
+  config['train_frames'] = 'frames_train_skirt.txt'
+  config['valid_frames'] = 'frames_valid_skirt.txt'
+else:
+  config['train_frames'] = 'frames_train.txt'
+  config['valid_frames'] = 'frames_valid.txt'
+
+config['smpl_res'] = 14475
+if config['is_cape']:
+  config['smpl_res'] = 3638
+
+config['latent_dim'] = 128
+
+config['bottleneck_scale'] = 2
+
+config['cond_ae'] = dict()
+config['cond_ae']['device'] = config['device']
+config['cond_ae']['is_label'] = False
+config['cond_ae']['latent_dim'] = config['latent_dim']
+config['cond_ae']['kernel_size'] = 7
+config['cond_ae']['bottleneck_size'] = 40 * config['bottleneck_scale']
+
+block_setting = [
+    CNBlockConfig(16, 32, 3),
+    CNBlockConfig(32, 64, 3),
+    CNBlockConfig(64, 128, 9),
+    CNBlockConfig(128, config['cond_ae']['bottleneck_size'], 9),
+    CNBlockConfig(config['cond_ae']['bottleneck_size'], None, 3),
+]
+
+block_setting_reverse = [
+    CNBlockConfig(config['cond_ae']['bottleneck_size'], 128, 3),
+    CNBlockConfig(128, 64, 9),
+    CNBlockConfig(64, 32, 9),
+    CNBlockConfig(32, 16, 3),
+    CNBlockConfig(16, 3, 3),
+    # CNBlockConfig(3, 3, 3),
+]
+
+config['cond_ae']['encoder_block_settings'] = block_setting
+config['cond_ae']['decoder_block_settings'] = block_setting_reverse
+config['cond_ae']['activation'] = 'tanh'
+
+
+config['lr'] = 0.0001
+config['cond_ae']['upscale_output'] = False
+
+config['signature'] = { 'id',
+                        'frame_nr',
+                        'uv_mask',
+                        'uv_normals',
+                        'mesh_mask',
+                        'mesh_normals',
+                        }

configs/iccvw/ae_static_1_skirt.py

@@ -0,0 +1,78 @@
+from utils.convnext_utils import CNBlockConfig
+
+config = dict()
+
+config['gpu'] = '3'
+config['name'] = 'static_iccvw_1_skirt'
+config['notes'] = ' ' 
+config['tags'] = ['torch', 'static'] 
+config['script'] = 'scripts/ae_static_train.py'
+
+config['pretrained_model'] = ''
+
+config['batch_size'] = 4
+config['train_dir'] = '/data-local/127/cloth3d_processed'
+config['valid_dir'] = '/data-local/127/cloth3d_processed'
+config['test_dir'] = '/data-net/dataset/test_ready'
+
+config['resolution'] = 512
+config['is_skirt'] = True
+config['dataset_in_memory'] = False
+config['is_cape'] = False
+config['remove_global_orientation'] = True
+config['device'] = 'cuda'
+
+if config['is_skirt']:
+  config['train_frames'] = 'frames_train_skirt.txt'
+  config['valid_frames'] = 'frames_valid_skirt.txt'
+else:
+  config['train_frames'] = 'frames_train.txt'
+  config['valid_frames'] = 'frames_valid.txt'
+
+config['smpl_res'] = 14475
+if config['is_cape']:
+  config['smpl_res'] = 3638
+
+config['latent_dim'] = 128
+
+config['bottleneck_scale'] = 2
+
+config['cond_ae'] = dict()
+config['cond_ae']['device'] = config['device']
+config['cond_ae']['is_label'] = False
+config['cond_ae']['latent_dim'] = config['latent_dim']
+config['cond_ae']['kernel_size'] = 7
+config['cond_ae']['bottleneck_size'] = 8 * config['bottleneck_scale']
+
+block_setting = [
+    CNBlockConfig(8, 16, 3),
+    CNBlockConfig(16, 32, 3),
+    CNBlockConfig(32, 64, 3),
+    CNBlockConfig(64, config['cond_ae']['bottleneck_size'], 9),
+    CNBlockConfig(config['cond_ae']['bottleneck_size'], None, 3),
+]
+
+block_setting_reverse = [
+    CNBlockConfig(config['cond_ae']['bottleneck_size'], 64, 3),
+    CNBlockConfig(64, 32, 9),
+    CNBlockConfig(32, 16, 3),
+    CNBlockConfig(16, 8, 3),
+    CNBlockConfig(8, 3, 3),
+    # CNBlockConfig(3, 3, 3),
+]
+
+config['cond_ae']['encoder_block_settings'] = block_setting
+config['cond_ae']['decoder_block_settings'] = block_setting_reverse
+config['cond_ae']['activation'] = 'tanh'
+
+
+config['lr'] = 0.0005
+config['cond_ae']['upscale_output'] = False
+
+config['signature'] = { 'id',
+                        'frame_nr',
+                        'uv_static',
+                        'mesh_static',
+                        'mesh_mask',
+                        'uv_mask',
+                        }

configs/iccvw/ae_static_5.py

@@ -0,0 +1,78 @@
+from utils.convnext_utils import CNBlockConfig
+
+config = dict()
+
+config['gpu'] = '5'
+config['name'] = 'static_iccvw_5'
+config['notes'] = 'lower learning rate' 
+config['tags'] = ['torch', 'static'] 
+config['script'] = 'scripts/ae_static_train.py'
+
+config['pretrained_model'] = 'checkpoints/iccvw/static_iccvw_3'
+
+config['batch_size'] = 4
+config['train_dir'] = '/data-local/127/cloth3d_processed'
+config['valid_dir'] = '/data-local/127/cloth3d_processed'
+config['test_dir'] = '/data-net/dataset/test_ready'
+
+config['resolution'] = 512
+config['is_skirt'] = False
+config['dataset_in_memory'] = False
+config['is_cape'] = False
+config['remove_global_orientation'] = True
+config['device'] = 'cuda'
+
+if config['is_skirt']:
+  config['train_frames'] = 'frames_train_skirt.txt'
+  config['valid_frames'] = 'frames_valid_skirt.txt'
+else:
+  config['train_frames'] = 'frames_train.txt'
+  config['valid_frames'] = 'frames_valid.txt'
+
+config['smpl_res'] = 14475
+if config['is_cape']:
+  config['smpl_res'] = 3638
+
+config['latent_dim'] = 128
+
+config['bottleneck_scale'] = 2
+
+config['cond_ae'] = dict()
+config['cond_ae']['device'] = config['device']
+config['cond_ae']['is_label'] = False
+config['cond_ae']['latent_dim'] = config['latent_dim']
+config['cond_ae']['kernel_size'] = 7
+config['cond_ae']['bottleneck_size'] = 8 * config['bottleneck_scale']
+
+block_setting = [
+    CNBlockConfig(8, 16, 3),
+    CNBlockConfig(16, 32, 3),
+    CNBlockConfig(32, 64, 3),
+    CNBlockConfig(64, config['cond_ae']['bottleneck_size'], 9),
+    CNBlockConfig(config['cond_ae']['bottleneck_size'], None, 3),
+]
+
+block_setting_reverse = [
+    CNBlockConfig(config['cond_ae']['bottleneck_size'], 64, 3),
+    CNBlockConfig(64, 32, 9),
+    CNBlockConfig(32, 16, 3),
+    CNBlockConfig(16, 8, 3),
+    CNBlockConfig(8, 3, 3),
+    # CNBlockConfig(3, 3, 3),
+]
+
+config['cond_ae']['encoder_block_settings'] = block_setting
+config['cond_ae']['decoder_block_settings'] = block_setting_reverse
+config['cond_ae']['activation'] = 'tanh'
+
+
+config['lr'] = 0.00005
+config['cond_ae']['upscale_output'] = False
+
+config['signature'] = { 'id',
+                        'frame_nr',
+                        'uv_static',
+                        'mesh_static',
+                        'mesh_mask',
+                        'uv_mask',
+                        }

custom_utils/custom_utils.py

@@ -9,9 +9,6 @@ from plotly.subplots import make_subplots
 from torchvision.transforms import Resize
 from torchvision.transforms.functional import InterpolationMode
 
-# from tensorflow_graphics.geometry.convolution.graph_pooling import pool, unpool
-# from tensorflow_graphics.geometry.convolution.graph_convolution import edge_convolution_template
-
 
 def m2uv_array(m2uv, is_smpl_orig=False):
     result = np.zeros((14475 if not is_smpl_orig else 3638, 2), dtype=np.int32)
@@ -74,35 +71,6 @@ class UV2Mesh(nn.Module):
         return x
 
 
-
-class UV2MeshLabel():
-    def __init__(self, res=256, is_skirt=False, is_smpl_orig=False):
-        super(UV2MeshLabel, self).__init__()
-        self.res = res
-        uv_map = np.load('custom_utils/stretched_uv_map_' + ('skirt' if is_skirt else '3') + ('_orig' if is_smpl_orig else '') + '.npy')
-        self.uv_map_discreet = np.floor(uv_map * self.res).astype(np.int32)
-        m2uv = np.load('custom_utils/m2uv' + ('_skirt' if is_skirt else '') + ('_orig' if is_smpl_orig else '') + '.npy', allow_pickle=True).item()
-        m2uv = m2uv_array(m2uv, is_smpl_orig=is_smpl_orig)
-        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, return_unfixed=False):
     V_mask = mask
     F_mask = np.take(V_mask, F.reshape(-1), axis=0)
@@ -119,20 +87,6 @@ def _maskFaces(V, F, mask, return_unfixed=False):
         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 = []
@@ -272,25 +226,6 @@ def plot_3d_collision(gt, collision_1, collision_2, collision_faces, collision_c
     return fig
 
 
-def batch_to_dict(batch, signature):
-    result = dict()
-    for i, key in enumerate(signature.keys()):
-            result[key] = batch[i]
-    return result
-
-conditions = ['pose_shape', 'uv_static', 'uv_mask', 'uv_shaped_body']
-def nest_conditions(batch):
-    batch['conditions'] = {d: batch.pop(d) for d in conditions}
-    return batch
-
-
-def sparse_to_tensor(L):
-    idx = np.mat([L.row, L.col]).transpose().astype(np.float32)
-    data = L.data.astype(dtype=np.float32)
-    shape = L.shape
-    return tf.SparseTensor(idx, data, shape)
-
-
 
 def get_summary_str(model):
     # from torchview import draw_graph
@@ -304,6 +239,7 @@ def get_summary_str(model):
 
     return [f'all_params: {total_params}, trainable: {trainable_params}']
 
+
 def get_summary_html(model):
     lines = get_summary_str(model)
     html = ''
@@ -312,88 +248,6 @@ def get_summary_html(model):
     return html
 
 
-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)
-
-
-def merge_verts(Vs, batchDim):
-    V = tf.concat(tf.unstack(Vs, num=batchDim), axis=0)
-    return V
-
-
-class GraphConvolution(nn.Module):
-    def __init__(self, shape, activation=None, name='gcn', L=None, batchSize=None, normalization=None):
-        super(GraphConvolution, self).__init__()
-        self.shape = shape
-        self.w = self.add_weight(name=name + '/weight', shape=(2, *shape), initializer=tf.initializers.glorot_normal, trainable=True)
-        self.b = self.add_weight(name=name + '/bias', shape=shape[-1], initializer=tf.initializers.zeros, trainable=True)
-        self.activation = activation or (lambda x: x)
-        self.normalization = normalization
-        self.batchSize = batchSize
-        self.indices = tf.cast(tf.linspace(0, self.batchSize * L.shape[1], self.batchSize+1), dtype=tf.int32)
-        self.L = repeat_laplacian(L, self.batchSize)
-
-    def gather(self):
-        return [self.w, self.b]
-
-    def _single(self, X):
-        X0 = tf.einsum('ab,bc->ac', X, self.w[0]) # Node
-        X1 = tf.einsum('ab,bc->ac', X, self.w[1]) # Neigh
-        # Neigh conv
-        X1 = edge_convolution_template(
-            X1, sparse_to_tensor(self.L), sizes=None,
-            edge_function=lambda x, y: y, 
-            reduction='weighted', 
-            edge_function_kwargs={}
-        ) 
-        X = X0 + X1 + self.b
-        return X
-
-    def call(self, X, training):
-        X = merge_verts(X, self.batchSize)
-        X = self._single(X)
-        X = tf.stack([X[self.indices[i-1]:self.indices[i]] for i in range(1, len(self.indices))])
-        X = self.activation(X)
-        if self.normalization:
-            X = self.normalization(X, training=training)
-        return X
-
-
-def replace_arrays(string):
-    out = string[:1]
-    i = 0
-    while i < len(string) and string[i+1:].find('[') != -1:
-        start = string[i+1:].find('[')
-        end = string[i+1+start:].find(']')
-        out += string[i+1:i+1+start]
-        array = string[i+1+start:i+1+start+end+1].replace('\n', '')
-        out += array
-        i = i + 1 + start + end
-
-    out += string[i+1:]
-
-    return out
-
-
-def group_inputs_for_levels(batch, config):
-    levels_input = []
-    for res in config['pyramid']['resolutions']:
-        res_modifier = '_' + str(res) if res != 512 else ''
-        if config['use_skinning']:
-            uv_orig = batch['uv_unposed' + res_modifier]
-        else:
-            uv_orig = batch['uv_posed' + res_modifier]
-        levels_input.append([uv_orig, batch['uv_static' + res_modifier], batch['uv_body_posed' + res_modifier], batch['uv_curvature' + res_modifier], batch['uv_mask' + res_modifier], batch['uv_mask']])
-    conds = [batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature']]
-    return levels_input, conds
-
-
 class Accumulator():
     def __init__(self):
         self.dict = dict()

data/downscale.py

@@ -0,0 +1,223 @@
+import argparse
+import os
+import sys
+import tqdm
+import cv2
+import torch
+import scipy.interpolate as spi
+import numpy as np
+import plotly.graph_objects as go
+import plotly.express as ex
+from torchvision.transforms import Resize
+from torchvision.transforms.functional import InterpolationMode
+
+os.environ["CUDA_VISIBLE_DEVICES"] = "4"
+
+sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
+from data.utils import Frame
+
+
+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)
+    # interpolator =  spi.RBFInterpolator(y=img_idx, d=img_data, neighbors=5, kernel='linear')
+
+    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_label(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), dtype=np.float32)
+    res[coords[:,0], coords[:,1]] = img_2
+    res[img_idx[:,0], img_idx[:,1]] = img_data
+
+    return res
+
+
+def generate_downscaled_uv_map(folder, frame: Frame, size):
+    img = np.load(os.path.join(folder, 'uv_mask', f'{frame.id}_{frame.garment}.npz')); img = img['uv_mask'].astype(np.float32)
+    down = resizers_nearest[size](torch.from_numpy(img[None,:,:]).to(device))
+    down = np.squeeze(down.cpu().numpy().astype(bool))
+    if not os.path.exists(os.path.join(folder, 'uv_mask' + '_' + str(size))):
+        os.makedirs(os.path.join(folder, 'uv_mask' + '_' + str(size)), exist_ok=True)
+    np.savez_compressed(os.path.join(folder, 'uv_mask' + '_' + str(size), f'{frame.id}_{frame.garment}.npz'), uv_mask = down)
+
+
+class UVMaskStore():
+    def __init__(self, folder, size):
+        self.store = dict()
+        self.folder = folder
+        self.size = size
+
+    def read(self, frame: Frame):
+        path = os.path.join(self.folder, 'uv_mask' + ('_' + str(self.size) if self.size != 512 else ''), f'{frame.id}_{frame.garment}.npz')
+        if not os.path.exists(path):
+            generate_downscaled_uv_map(self.folder, frame, self.size)
+
+        mask = np.load(path)
+        mask = mask['uv_mask']
+        return mask  
+
+    def __getitem__(self, frame: Frame) -> dict:
+        if (frame.id, frame.garment) not in self.store:
+            self.store[(frame.id, frame.garment)] = self.read(frame)
+
+        return self.store[(frame.id, frame.garment)]
+
+
+def downscale_uv_maps_posed_body(folder, type, sizes):
+    for size in sizes:
+        os.makedirs(os.path.join(folder, type + str(size)), exist_ok=True)
+
+    frames = sorted(os.listdir(os.path.join(folder, type)))[int(args.start):int(args.end)]
+    frames = [Frame(frame.split('_')[0], frame.split('_')[1].split('.')[0], '') for frame in frames]
+
+    for frame in tqdm.tqdm(frames):
+        img = np.load(os.path.join(folder, type, f'{frame.id}_{frame.frame_nr}.npz'))['uv_body_posed']
+        img = fill(img, uv_masks_body[512])
+        for size in sizes:
+            mask = uv_masks_body[size]
+            down = resizers[size](torch.from_numpy(img.transpose((2,0,1))).to(device)).cpu().numpy()
+            down = down * mask[None,:,:]
+            np.savez_compressed(os.path.join(folder, type + str(size), f'{frame.id}_{frame.frame_nr}.npz'), uv_body_posed = down)
+
+
+def downscale_uv_maps_posed_no_orient(folder, type, sizes):
+    for size in sizes:
+        os.makedirs(os.path.join(folder, type + '_' + str(size)), exist_ok=True)
+
+    frames = sorted(os.listdir(os.path.join(folder, type)))[int(args.start):int(args.end)]
+    frames = [Frame(frame.split('_')[0], frame.split('_')[1], frame.split('_')[2].split('.')[0]) for frame in frames]
+
+    for frame in tqdm.tqdm(frames):
+        img = np.load(os.path.join(folder, type, f'{frame.id}_{frame.frame_nr}_{frame.garment}.npz'))['uv_posed']
+        img = fill(img, uv_mask_store[512][frame])
+        for size in sizes:
+            mask = uv_mask_store[size][frame]
+            down = resizers[size](torch.from_numpy(img.transpose((2,0,1))).to(device)).cpu().numpy()
+            down = down * mask[None,:,:]
+            np.savez_compressed(os.path.join(folder, type + '_' + str(size), f'{frame.id}_{frame.frame_nr}_{frame.garment}.npz'), uv_posed = down)
+
+
+def downscale_uv_maps_normals(folder, type, sizes):
+    for size in sizes:
+        os.makedirs(os.path.join(folder, type + '_' + str(size)), exist_ok=True)
+
+    frames = sorted(os.listdir(os.path.join(folder, type)))[int(args.start):int(args.end)]
+    frames = [Frame(frame.split('_')[0], frame.split('_')[1], frame.split('_')[2].split('.')[0]) for frame in frames]
+
+    for frame in tqdm.tqdm(frames):
+        img = np.load(os.path.join(folder, type, f'{frame.id}_{frame.frame_nr}_{frame.garment}.npz'))['uv_normals']
+        img = fill(img, uv_mask_store[512][frame])
+        for size in sizes:
+            mask = uv_mask_store[size][frame]
+            down = resizers[size](torch.from_numpy(img.transpose((2,0,1))).to(device)).cpu().numpy()
+            down = down * mask[None,:,:]
+            np.savez_compressed(os.path.join(folder, type + '_' + str(size), f'{frame.id}_{frame.frame_nr}_{frame.garment}.npz'), uv_normals = down)
+
+
+def downscale_uv_maps_unposed(folder, type, sizes):
+    for size in sizes:
+        os.makedirs(os.path.join(folder, type + '_' + str(size)), exist_ok=True)
+
+    frames = sorted(os.listdir(os.path.join(folder, type)))[int(args.start):int(args.end)]
+    frames = [Frame(frame.split('_')[0], frame.split('_')[1], frame.split('_')[2].split('.')[0]) for frame in frames]
+
+    for frame in tqdm.tqdm(frames):
+        img = np.load(os.path.join(folder, type, f'{frame.id}_{frame.frame_nr}_{frame.garment}.npz'))['uv_unposed']
+        img = fill(img, uv_mask_store[512][frame])
+        for size in sizes:
+            mask = uv_mask_store[size][frame]
+            down = resizers[size](torch.from_numpy(img.transpose((2,0,1))).to(device)).cpu().numpy()
+            down = down * mask[None,:,:]
+            np.savez_compressed(os.path.join(folder, type + '_' + str(size), f'{frame.id}_{frame.frame_nr}_{frame.garment}.npz'), uv_unposed = down)
+
+
+def downscale_uv_maps_static(folder, type, sizes):
+    for size in sizes:
+        os.makedirs(os.path.join(folder, type + '_' + str(size)), exist_ok=True)
+
+    frames = sorted(os.listdir(os.path.join(folder, type)))
+    import random
+    random.shuffle(frames)
+    frames = [(frame.split('_')[0], frame.split('_')[1].split('.')[0]) for frame in frames]
+
+    for frame in tqdm.tqdm(frames):
+        img = np.load(os.path.join(folder, type, f'{frame[0]}_{frame[1]}.npz'))['uv_static']
+        img = fill(img, uv_mask_store[512][Frame(frame[0], 0, frame[1])])
+        for size in sizes:
+            mask = uv_mask_store[size][Frame(frame[0], 0, frame[1])]
+            down = resizers[size](torch.from_numpy(img.transpose((2,0,1))).to(device)).cpu().numpy()
+            down = down * mask[None,:,:]
+            np.savez_compressed(os.path.join(folder, type + '_' + str(size), f'{frame[0]}_{frame[1]}.npz'), uv_static = down)
+
+
+def downscale_uv_maps_mask(folder, type, size):
+    os.makedirs(os.path.join(folder, type + '_' + str(size)), exist_ok=True)
+
+    frames = sorted(os.listdir(os.path.join(folder, type)))[int(args.start):int(args.end)]
+
+    for frame in tqdm.tqdm(frames):
+        id, _ = frame.split('.')
+        img = np.load(os.path.join(folder, type, frame)); img = img['uv_mask'].astype(np.float32)
+        down = tf.image.resize(img[:,:,tf.newaxis], size=(size, int(size/2)), method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)
+        down = np.squeeze(down.numpy().astype(bool))
+        np.savez_compressed(os.path.join(folder, type + '_' + str(size), frame), uv_mask = down)
+
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument("-f", "--folder", help = "end of range", default='/data-local/127/cloth3d_processed')
+    parser.add_argument("-s", "--start", help = "start of range", default=0)
+    parser.add_argument("-e", "--end", help = "end of range", default=25000)
+    parser.add_argument("-n", "--name", help = "end of range", default='lhunor')
+    parser.add_argument("-d", "--debug", help = "end of range", default=False, action='store_true')
+    args = parser.parse_args()
+
+    # sizes = [32, 64, 128, 256]
+    device = 'cuda'
+    sizes = [32, 64, 128, 256]
+    folder = args.folder
+    resizers = { res: Resize(size=(res, int(res/2)), interpolation=InterpolationMode.BICUBIC).to(device) for res in sizes }
+    resizers_nearest = { res: Resize(size=(res, int(res/2)), interpolation=InterpolationMode.NEAREST_EXACT).to(device) for res in sizes }
+
+    uv_masks_body = dict()
+    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)
+    uv_masks_body[512] = uv_mask_body.astype(bool)
+    for res in [32, 64, 128, 256]:
+        down = resizers_nearest[res](torch.from_numpy(uv_mask_body[None,:,:].astype(np.int32)).to(device))
+        down = np.squeeze(down.cpu().numpy().astype(bool))
+        uv_masks_body[res] = down
+        
+
+    uv_mask_store = { size: UVMaskStore(folder, size=size) for size in sizes }
+    uv_mask_store[512] = UVMaskStore(folder, size=512)
+
+    # downscale_uv_maps_mask(folder=folder, type='uv_mask', size=32)
+    # downscale_uv_maps_mask(folder=folder, type='uv_mask', size=64)
+    # downscale_uv_maps_mask(folder=folder, type='uv_mask', size=128)
+    # downscale_uv_maps_mask(folder=folder, type='uv_mask', size=256)
+    # downscale_uv_maps_static(folder=folder, type='uv_static', sizes=sizes)
+
+    downscale_uv_maps_unposed(folder=folder, type='uv_unposed_no_shape', sizes=sizes)
+    downscale_uv_maps_posed_no_orient(folder=folder, type='uv_posed_no_orient', sizes=sizes)
+    downscale_uv_maps_normals(folder=folder, type='uv_normals', sizes=sizes)
+    downscale_uv_maps_posed_body(folder=folder, type='uv_body_posed_no_orient', sizes=sizes)
+
+
+    print()

data/preprocessing.py

@@ -30,7 +30,7 @@ import torch
 from scipy.spatial.transform import Rotation as R
 
 sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
-from data.cloth3d_io import readPC2, loadInfo, writePC2, readOBJ, quads2tris
+from data.cloth3d_io import readPC2, loadInfo, writePC2, readOBJ, quads2tris, readPC2Frame
 from data.utils import Frame
 from smpl.smpl_torch import SMPL
 from custom_utils.custom_utils import _maskFaces
@@ -160,7 +160,7 @@ def deorient(verts, pose):
 
 
 def create_mesh_posed_no_orient_and_mask(folder, folder_reg, out_folder):
-    ids = os.listdir(folder)
+    ids = sorted(os.listdir(folder))
     os.makedirs(os.path.join(out_folder, 'mesh_posed_no_orient'), exist_ok=True)
     os.makedirs(os.path.join(out_folder, 'mesh_mask'), exist_ok=True)
 
@@ -179,12 +179,18 @@ def create_mesh_posed_no_orient_and_mask(folder, folder_reg, out_folder):
                 continue
         
             verts = readPC2(os.path.join(garment_dir, garment + '.pc16'), True)['V']
+            # if info['poses'].shape[1] == verts.shape[0]:
+            #     continue
             reg_data = sio.loadmat(os.path.join(folder_reg, id, garment + '_weights_fixed.mat'))
             verts = (verts[:, reg_data['F'], :] * reg_data['W'][None, :, :, None]).sum(axis=2)
             mask_idx = sio.loadmat(os.path.join(folder_reg, id, garment + '_mask.mat'))['mask_idx'].reshape((-1)) - 1
             mask = np.zeros((14475), dtype=bool)
             mask[mask_idx] = True
 
+            if info['poses'].shape[1] != verts.shape[0]:
+                assert False, 'Somethign wrong here'
+                info['poses'] = info['poses'][:, :verts.shape[0]]
+
             verts = deorient(verts, info['poses'])
 
             # # plot as scatter plot with plotly
@@ -193,13 +199,12 @@ def create_mesh_posed_no_orient_and_mask(folder, folder_reg, out_folder):
             # fig = go.Figure(data=[go.Scatter3d(x=verts[frame,:,0], y=verts[frame,:,2], z=verts[frame,:,1], mode='markers', marker=dict(size=1))])
             # fig.show()
 
-
             writePC2(os.path.join(out_folder, 'mesh_posed_no_orient', f'{id}_{garment}.pc16'), verts, True)
             np.save(os.path.join(out_folder, 'mesh_mask', f'{id}_{garment}.npy'), mask)
 
 
 def create_pose_shape(folder, out_folder):
-    ids = os.listdir(folder)
+    ids = sorted(os.listdir(folder))
     os.makedirs(os.path.join(out_folder, 'pose'), exist_ok=True)
     os.makedirs(os.path.join(out_folder, 'shape'), exist_ok=True)
 
@@ -216,7 +221,7 @@ def create_pose_shape(folder, out_folder):
 
 
 def create_mesh_static(folder, out_folder):
-    ids = os.listdir(folder)
+    ids = sorted(os.listdir(folder))
     os.makedirs(os.path.join(out_folder, 'mesh_static'), exist_ok=True)
 
     for id in tqdm(ids):
@@ -247,7 +252,7 @@ def create_mesh_static(folder, out_folder):
 
 
 def create_mesh_body_posed_no_orient(folder, out_folder):
-    ids = os.listdir(folder)
+    ids = sorted(os.listdir(folder))
     os.makedirs(os.path.join(out_folder, 'mesh_body_posed_no_orient'), exist_ok=True)
 
     for id in tqdm(ids):
@@ -288,10 +293,11 @@ def calculate_normals(vertices, faces):
 
 
 def create_mesh_normals(folder, out_folder):
-    ids = os.listdir(folder)
+    ids = sorted(os.listdir(folder))
     os.makedirs(os.path.join(out_folder, 'mesh_normals'), exist_ok=True)
     # F = np.load('smpl/smpl_faces.npy')
     # F_skirt = np.load('smpl/smpl_faces_skirt.npy')
+    # ids = ['00438']
 
     for id in tqdm(ids):
         garment_dir = os.path.join(folder, id)
@@ -306,7 +312,7 @@ def create_mesh_normals(folder, out_folder):
                 continue
 
             verts = readPC2(os.path.join(garment_dir, garment + '.pc16'), True)['V']
-            verts = deorient(verts, info['poses'])
+            verts = deorient(verts, info['poses'][:, :verts.shape[0]])
             _, faces, _, _ = readOBJ(os.path.join(garment_dir, garment + '.obj'))
             faces = quads2tris(faces)
             normals = calculate_normals(verts, faces.astype(np.int32))
@@ -413,7 +419,7 @@ def create_mesh_normals_cape(folder, out_folder):
 
 
 def verify_registration_ans_mask(folder, out_folder):
-    ids = os.listdir(folder)
+    ids = sorted(os.listdir(folder))
 
     for id in ids:
         garment_dir = os.path.join(folder, id)
@@ -621,9 +627,80 @@ def create_frames_file(folder, out_folder, split=0.9):
                 f.write(f'{frame.id} {frame.frame_nr} {frame.garment}\n')
 
 
+def create_frames_file_partial(folder, out_folder, split=0.9, part=0.1):
+    ids = os.listdir(folder)
+
+    frames_non_skirt = []
+    frames_skirt = []
+    ids_non_skirt = set()
+    ids_skirt = set()
+
+    for id in tqdm(ids):
+        garment_dir = os.path.join(folder, id)
+        if not os.path.isdir(garment_dir):
+            continue
+
+        info = sio.loadmat(os.path.join(folder, id, 'info.mat'))
+        garments = [f.replace('.obj','') for f in os.listdir(garment_dir) if f.endswith('.obj')]
+        frames = info['poses'].shape[1]
+
+        assert info['poses'].shape[0] == 72
+
+        for garment in garments:
+            for i in range(frames):
+                if garment == 'Skirt' or garment == 'Dress':
+                    frames_skirt.append(Frame(id, i, garment))
+                    ids_skirt.add(id)
+                else:
+                    frames_non_skirt.append(Frame(id, i, garment))
+                    ids_non_skirt.add(id)
+
+    nr_train_non_skirt = int(len(ids_non_skirt) * split * part)
+    nr_train_skirt = int(len(ids_skirt) * split * part)
+
+    ids_skirt = list(ids_skirt)
+    ids_non_skirt = list(ids_non_skirt)
+
+    with open(os.path.join(out_folder, f'frames_train_{int(part*100)}p.txt'), 'w') as f:
+        for frame in frames_non_skirt:
+            if frame.id in ids_non_skirt[:nr_train_non_skirt]:
+                f.write(f'{frame.id} {frame.frame_nr} {frame.garment}\n')
+
+    with open(os.path.join(out_folder, f'frames_train_skirt_{int(part*100)}p.txt'), 'w') as f:
+        for frame in frames_skirt:
+            if frame.id in ids_skirt[:nr_train_skirt]:
+                f.write(f'{frame.id} {frame.frame_nr} {frame.garment}\n')
+
+
+
+def check_irregular_poses(folder, folder_reg, out_folder):
+    frame: Frame = Frame('01512', 50, 'Dress')
+    mesh_posed = readPC2Frame(os.path.join(out_folder, 'mesh_posed_no_orient', f'{frame.id}_{frame.garment}.pc16'), frame.frame_nr, True)
+    # mesh_static = np.load(f'{out_folder}/mesh_static/{frame.id}_{frame.garment}.npy')
+
+    mesh_mask = np.load(f'{out_folder}/mesh_mask/{frame.id}_{frame.garment}.npy')
+    mesh_body_posed = readPC2Frame(os.path.join(out_folder, 'mesh_body_posed_no_orient', f'{frame.id}.pc16'), frame.frame_nr, True)
+
+    F = np.load('smpl/smpl_faces_skirt.npy')
+    _, faces = _maskFaces(np.zeros((14475, 3)), F, mesh_mask)
+
+    fig = go.Figure(data=[go.Mesh3d(x=mesh_posed[:,0], y=mesh_posed[:,2], z=mesh_posed[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='lightpink', opacity=1)])
+    fig.add_trace(go.Mesh3d(x=mesh_body_posed[:,0], y=mesh_body_posed[:,2], z=mesh_body_posed[:,1], i=F[:,0], j=F[:,1], k=F[:,2], color='lightblue', opacity=1))
+    fig.update_layout(scene = dict(
+        xaxis = dict(range=[-1,1],),
+                    yaxis = dict(range=[-1,1],),
+                    zaxis = dict(range=[-1,1],),),)
+    fig.update_layout(scene_aspectmode='cube')
+    fig.show()
+    pass
+
 if __name__ == '__main__':
 
 
+    # folder = '/data-net/dataset/cloth3d/test'
+    # folder_reg = '/data-local/127/cloth3d_registration/test'
+    # out_folder = '/data-local/127/cloth3d_processed_test'
+
     folder = '/data-net/dataset/cloth3d/train'
     folder_reg = '/data-local/127/cloth3d_registration/train'
     out_folder = '/data-local/127/cloth3d_processed'
@@ -635,6 +712,8 @@ if __name__ == '__main__':
     # split_data(folder, 0.9)
     # create_extras_file(folder)
 
+    # check_irregular_poses(folder, folder_reg, out_folder)
+
     # create_mesh_posed_no_orient_and_mask(folder, folder_reg, out_folder)
     # verify_registration_ans_mask(folder, out_folder)
     # create_pose_shape(folder, out_folder)
@@ -644,6 +723,7 @@ if __name__ == '__main__':
     # create_mesh_unposed_no_shape(folder, folder_reg, out_folder, start=args.start, end=args.end)
 
     # create_frames_file(folder, out_folder)
+    create_frames_file_partial(folder, out_folder, part=0.5)
     
-    folder = '/data-local/127_2/cape_processed/test/female'
+    # folder = '/data-local/127_2/cape_processed/test/female'
-    create_mesh_normals_cape(folder, None)
+    # create_mesh_normals_cape(folder, None)

scripts/ae_body_train.py

@@ -11,10 +11,12 @@ parser.add_argument('--config', type=str)
 parser.add_argument('--debug', action='store_true')
 args = vars(parser.parse_args())
 
-config_module = importlib.import_module(args['config'])
-config = config_module.config
 if args['debug']:
     from configs.iccvw.ae_body_cape_m import config
+else:
+    config_module = importlib.import_module(args['config'])
+    config = config_module.config
+
 os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
 os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
 if args['debug']:
@@ -32,10 +34,7 @@ import plotly.graph_objects as go
 import plotly.express as ex
 import torch.optim as optim
 
-if config['is_skirt']:
+from losses.losses import edge_loss
-    from losses.losses import edge_loss_skirt as edge_loss
-else:
-    from losses.losses import edge_loss
 
 from data.dataset import Dataset
 # from torch.utils.data import DataLoader
@@ -65,7 +64,7 @@ def init_wandb():
             artifact.add_file(file_name)
     run.log_artifact(artifact)
 
-uv2mesh = UV2Mesh(res=256, is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape'])
+uv2mesh = UV2Mesh(res=256, is_skirt=False, is_smpl_orig=config['is_cape'])
 
 
 
@@ -122,7 +121,7 @@ def plot_everything(batch: dict, type: str):
     img = torch.concat((orig_img, pred_img), axis=2)
     imgs = [wandb.Image((img[i].cpu().numpy() + 1) / 2.0) for i in range(4)]
 
-    F = np.load("smpl/smpl_faces" + skirt_modifier + ("_orig" if config['is_cape'] else "") + ".npy")
+    F = np.load("smpl/smpl_faces.npy")
     garment_unposed, faces = mesh_pred.cpu().numpy(), F
     gt = [mesh_orig[i].cpu().numpy() for i in range(pred.shape[0])]
     
@@ -176,9 +175,7 @@ if __name__ == '__main__':
     if not args['debug']:
         init_wandb()
 
-    skirt_modifier = '_skirt' if config['is_skirt'] else ''
+    uv_mask_body = np.load('custom_utils/mask_256x512.npy')
-
-    uv_mask_body = np.load('custom_utils/mask_256x512' + ('_skirt' if config['is_skirt'] else '') + '.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)
     uv_mask_body = torch.from_numpy(uv_mask_body).float().to(config['device'])
 

scripts/ae_static_train.py

@@ -94,7 +94,7 @@ def train_step(batch):
         mesh_orig = batch['mesh_static']
         l1_loss = torch.mean(torch.sum((torch.abs(batch['uv_static'] - pred)), axis=1))
         e_loss = 1e-2 * torch.mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']))
-        metric = torch.sum(torch.sqrt(torch.sum(batch['mesh_mask'][:,:,None] * (mesh_orig - mesh_pred.detach()) ** 2, axis=-1))) / torch.sum(batch['mesh_mask'])
+        metric = torch.sum(torch.sqrt(torch.sum(batch['mesh_mask'][:,:,None] * (mesh_orig - mesh_pred) ** 2, axis=-1))) / torch.sum(batch['mesh_mask'])
     
     return { 'train/loss': loss, 'train/edge_loss': e_loss, 'train/l1_loss': l1_loss, 'train/euclid_error': metric, 'train/ms_ssim': ms_ssim, }
         

scripts/upscaler_train.py

@@ -11,7 +11,7 @@ parser.add_argument('--debug', action='store_true')
 args = vars(parser.parse_args())
 
 os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
-os.environ["CUDA_VISIBLE_DEVICES"] = '6'
+os.environ["CUDA_VISIBLE_DEVICES"] = '2'
 
 
 import time
@@ -49,26 +49,36 @@ gender_modifier = '' if not is_cape else ('_m' if is_male else '_f')
 smpl_res = 3638 if is_cape else 14475
 use_cache = False
 name=f'iccv_{low_res}' + skirt_modifier + cape_modifier + gender_modifier
+train_dir = '/data-local/127/cloth3d_processed'
+valid_dir = '/data-local/127/cloth3d_processed'
 
-if not is_skirt:
+if is_skirt:
-    train_dir = '/data-local/127/train'
+  train_frames = 'frames_train_skirt.txt'
-    valid_dir = '/data-local/127/valid'
+  valid_frames = 'frames_valid_skirt.txt'
 else:
-    train_dir = '/data-local/128/train_skirt'
+  train_frames = 'frames_train.txt'
-    valid_dir = '/data-local/127/valid_skirt'
+  valid_frames = 'frames_valid.txt'
+
+if is_cape:
+    train_frames = None
+    valid_frames = None
 
 if is_cape:
     if is_male:
-        train_dir = '/data-local/127/cape_processed/train/male'
+        train_dir = '/data-local/127_2/cape_processed/train/male'
-        valid_dir = '/data-local/127/cape_processed/test/male'
+        valid_dir = '/data-local/127_2/cape_processed/test/male'
     else:
-        train_dir = '/data-local/127/cape_processed/train/female'
+        train_dir = '/data-local/127_2/cape_processed/train/female'
-        valid_dir = '/data-local/127/cape_processed/test/female'   
+        valid_dir = '/data-local/127_2/cape_processed/test/female'   
 
+print(f'nohup python scripts/upscaler_train.py --name {name} > nohup/up/up_{name}.out 2>&1 &')
+if args['debug']:
+    exit()
 
 
 def init_wandb():
-    run = wandb.init(project='upscale', name=name, notes='', dir='/data-local/127/temp', tags=['upscaler', 'iccv'])
+    wandb.login(host='https://wandb2.lnrweb.com', key='local-2a78b434e9bc51e9f12625f8321a9e3ee34af16d', relogin=True, force=True)
+    run = wandb.init(project='upscale', name=name, notes='', dir='/data-local/127/temp', tags=['upscaler', 'iccvw'])
     artifact = wandb.Artifact('code', type='code')
     artifact.add_file('scripts/upscaler_train.py')
     artifact.add_file('losses/losses.py')
@@ -255,8 +265,8 @@ if __name__ == '__main__':
                         }
     
     # dataset = Dataset('/data-local/127/train', config['signature'], resolutions=[512], device=device, frames_file='frames_jumpsuit.txt')
-    dataset_train = Dataset(train_dir, device='cuda', use_cache=use_cache, signature=signature, frames_file=None, resolutions=[512, low_res], is_cape=is_cape)
+    dataset_train = Dataset(train_dir, device='cuda', use_cache=use_cache, signature=signature, frames_file=train_frames, resolutions=[512, low_res], is_cape=is_cape)
-    dataset_valid = Dataset(valid_dir, device='cuda', use_cache=use_cache, signature=signature, frames_file=None, resolutions=[512, low_res], is_cape=is_cape)
+    dataset_valid = Dataset(valid_dir, device='cuda', use_cache=use_cache, signature=signature, frames_file=valid_frames, resolutions=[512, low_res], is_cape=is_cape)
     train_dataloader = DataLoader(dataset_train, batch_size=4, drop_last=True, shuffle=True)
     valid_dataloader = DataLoader(dataset_valid, batch_size=4, drop_last=True, shuffle=True)
     train_accumulator = Accumulator()
@@ -279,13 +289,13 @@ if __name__ == '__main__':
     _MSSSIM_small = MS_SSIM(data_range=2, size_average=True, channel=3, weights=[1.0])
 
     step = 0
-    ckpt_saver = CheckpointSaver('checkpoints/iccv/torch/' + name, top_n=5)
+    ckpt_saver = CheckpointSaver('checkpoints/iccvw/' + name, top_n=5)
 
     print(f'started training: {name}')
 
     if not args['debug']:
         # model = torch.compile(model)
-        wandb.watch(model, log='all', log_freq=1000, log_graph=True)
+        wandb.watch(model, log='all', log_freq=10000, log_graph=True)
 
     train()
 

utils/generate_figures.py

@@ -0,0 +1,126 @@
+import os
+import sys
+sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
+
+import numpy as np
+
+from custom_utils.custom_utils import save_stl, write_obj
+from data.cloth3d_io import readPC2Frame
+from data.utils import Frame
+from custom_utils.custom_utils import _maskFaces
+import plotly.graph_objects as go
+from plotly.subplots import make_subplots
+
+
+def remove_long_edges(vertices, faces):
+    # Compute edge lengths
+    edges = np.vstack((faces[:,[0,1]], faces[:,[1,2]], faces[:,[2,0]]))
+    edge_lengths = np.linalg.norm(vertices[edges[:,0]] - vertices[edges[:,1]], axis=1)
+
+    # Find edges longer than 0.1
+    long_edges = edges[edge_lengths > 0.1]
+
+    # Find vertices that are part of long edges
+    vertices_to_remove = np.unique(long_edges)
+
+    # Remove vertices and update faces
+    mask = np.ones(vertices.shape[0], dtype=bool)
+    mask[vertices_to_remove] = False
+    new_vertices = vertices[mask]
+
+
+    return new_vertices, mask
+
+
+def plot_unposing_example():
+    frame = Frame('00438', 146, 'Jumpsuit')
+    mesh_static = np.load(f'{folder_test}/mesh_static/{frame.id}_{frame.garment}.npy')
+
+    mesh_mask = np.load(f'{folder_test}/mesh_mask/{frame.id}_{frame.garment}.npy')
+    mesh_posed = readPC2Frame(os.path.join(folder_test, 'mesh_posed_no_orient', f'{frame.id}_{frame.garment}.pc16'), frame.frame_nr, True)
+    mesh_unposed = readPC2Frame(os.path.join(folder_test, 'mesh_unposed_no_shape', f'{frame.id}_{frame.garment}.pc16'), frame.frame_nr, True)
+    mesh_unposed_shaped = readPC2Frame(os.path.join('temp', f'{frame.id}_{frame.garment}_unposed_shaped.pc16'), frame.frame_nr, True)[mesh_mask]
+
+
+    F = np.load('smpl/smpl_faces.npy')
+    _, faces = _maskFaces(np.zeros((14475, 3)), F, mesh_mask)
+
+    mesh_static, v_mask = remove_long_edges(mesh_static, faces.astype(np.int32))
+
+    mesh_posed = mesh_posed[v_mask]
+    mesh_unposed = mesh_unposed[v_mask]
+    mesh_unposed_shaped = mesh_unposed_shaped[v_mask]
+    new_mask = np.zeros_like(mesh_mask)
+    new_mask[mesh_mask] = v_mask
+    _, faces = _maskFaces(np.zeros((14475, 3)), F, new_mask)
+
+    fig = make_subplots(rows=1, cols=4, specs=[[{'type': 'mesh3d'}, {'type': 'mesh3d'}, {'type': 'mesh3d'}, {'type': 'mesh3d'}]])
+
+    fig.add_trace(go.Mesh3d(x=mesh_static[:,0], y=mesh_static[:,2], z=mesh_static[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='lightpink', opacity=1), row=1, col=1)
+    fig.add_trace(go.Mesh3d(x=mesh_posed[:,0], y=mesh_posed[:,2], z=mesh_posed[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='lightblue', opacity=1), row=1, col=2)
+    fig.add_trace(go.Mesh3d(x=mesh_unposed_shaped[:,0], y=mesh_unposed_shaped[:,2], z=mesh_unposed_shaped[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='lightgreen', opacity=1), row=1, col=3)
+    fig.add_trace(go.Mesh3d(x=mesh_unposed[:,0], y=mesh_unposed[:,2], z=mesh_unposed[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='lightyellow', opacity=1), row=1, col=4)
+
+    fig.update_layout(scene=dict(
+        xaxis=dict(range=[-1, 1]),
+        yaxis=dict(range=[-1, 1]),
+        zaxis=dict(range=[-1, 1]),
+    ), scene_aspectmode='cube')
+
+    # fig.show()
+    faces = faces.astype(np.uint32)
+
+    save_stl(f'temp/unposing_plots/{frame.id}_{frame.frame_nr}_{frame.garment}_mesh_static.stl', mesh_static, faces)
+    save_stl(f'temp/unposing_plots/{frame.id}_{frame.frame_nr}_{frame.garment}_mesh_posed.stl', mesh_posed, faces)
+    save_stl(f'temp/unposing_plots/{frame.id}_{frame.frame_nr}_{frame.garment}_mesh_unposed.stl', mesh_unposed, faces)
+    save_stl(f'temp/unposing_plots/{frame.id}_{frame.frame_nr}_{frame.garment}_mesh_unposed_shaped.stl', mesh_unposed_shaped, faces)
+
+    pass
+
+
+def plot_normals_example():
+    frame = Frame('00438', 146, 'Jumpsuit')
+    mesh_static = np.load(f'{folder_test}/mesh_static/{frame.id}_{frame.garment}.npy')
+
+    mesh_mask = np.load(f'{folder_test}/mesh_mask/{frame.id}_{frame.garment}.npy')
+    mesh_posed = readPC2Frame(os.path.join(folder_test, 'mesh_posed_no_orient', f'{frame.id}_{frame.garment}.pc16'), frame.frame_nr, True)
+    # mesh_unposed = readPC2Frame(os.path.join(folder_test, 'mesh_unposed_no_shape', f'{frame.id}_{frame.garment}.pc16'), frame.frame_nr, True)
+    # mesh_unposed_shaped = readPC2Frame(os.path.join('temp', f'{frame.id}_{frame.garment}_unposed_shaped.pc16'), frame.frame_nr, True)[mesh_mask]
+    mesh_normals = readPC2Frame(os.path.join(folder_test, 'mesh_normals', f'{frame.id}_{frame.garment}.pc16'), frame.frame_nr, True)
+    uv_normals = np.load(os.path.join(folder_test, 'uv_normals', f'{frame.id}_{frame.frame_nr}_{frame.garment}.npz'))
+
+    F = np.load('smpl/smpl_faces.npy')
+    _, faces = _maskFaces(np.zeros((14475, 3)), F, mesh_mask)
+
+    mesh_static, v_mask = remove_long_edges(mesh_static, faces.astype(np.int32))
+
+    mesh_posed = mesh_posed[v_mask]
+    # mesh_unposed = mesh_unposed[v_mask]
+    # mesh_unposed_shaped = mesh_unposed_shaped[v_mask]
+    new_mask = np.zeros_like(mesh_mask)
+    new_mask[mesh_mask] = v_mask
+    _, faces = _maskFaces(np.zeros((14475, 3)), F, new_mask)
+
+    fig = make_subplots(rows=1, cols=4, specs=[[{'type': 'mesh3d'}, {'type': 'mesh3d'}, {'type': 'mesh3d'}, {'type': 'mesh3d'}]])
+
+    fig.add_trace(go.Mesh3d(x=mesh_static[:,0], y=mesh_static[:,2], z=mesh_static[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='lightpink', opacity=1), row=1, col=1)
+    fig.add_trace(go.Mesh3d(x=mesh_posed[:,0], y=mesh_posed[:,2], z=mesh_posed[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='lightblue', opacity=1), row=1, col=2)
+    fig.add_trace(go.Mesh3d(x=mesh_unposed_shaped[:,0], y=mesh_unposed_shaped[:,2], z=mesh_unposed_shaped[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='lightgreen', opacity=1), row=1, col=3)
+    fig.add_trace(go.Mesh3d(x=mesh_unposed[:,0], y=mesh_unposed[:,2], z=mesh_unposed[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='lightyellow', opacity=1), row=1, col=4)
+
+    fig.update_layout(scene=dict(
+        xaxis=dict(range=[-1, 1]),
+        yaxis=dict(range=[-1, 1]),
+        zaxis=dict(range=[-1, 1]),
+    ), scene_aspectmode='cube')
+
+
+
+if __name__ == '__main__':
+    folder_train = '/data-local/127/cloth3d_processed'
+    folder_test = '/data-local/127/cloth3d_processed_test'
+    
+    # plot_unposing_example()
+    plot_normals_example()
+
+    pass