making ready for skirts

auto_encoder_body.py

@@ -142,14 +142,15 @@ if __name__ == '__main__':
     parser.add_argument('--name', type=str)
     args = vars(parser.parse_args())
 
-    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)
 
     # TODO: different resolutions not working yet, model should be ready though
     resolution_modifier = ''
     if config['resolution'] != 512:
         resolution_modifier = '_' + str(config['resolution'])
 
+    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()
 
     dataset_train = factory.get_dataset(folder=config['train_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'], use_filled=config['use_filled'], resolutions=[config['resolution']], in_memory=config['dataset_in_memory'])

auto_encoder_curv.py

@@ -22,7 +22,7 @@ from models.cond_encoder import AutoEncoder
 testing = False
 def init_wandb():
     if not testing:
-        run = wandb.init(project='vae', name=config['name'], notes='convnext + convnext, 16*8*32, ce loss. flattened latent dim of 128. trained on 100 garments. same as 7')
+        run = wandb.init(project='vae', name=config['name'], notes=config['notes'])
         artifact = wandb.Artifact('code', type='code')
         artifact.add_file('auto_encoder_curv.py')
         artifact.add_file('losses.py')
@@ -79,7 +79,7 @@ def train_step(batch):
     with tf.GradientTape() as tape:
         pred = model([batch['uv_curvature'], batch['uv_mask']], training=True)
 
-        ce_uv_loss = cce(batch['uv_curvature'], pred)
+        ce_uv_loss = scce(tf.math.argmax(batch['uv_curvature'], axis=-1), pred)
         ms_ssim = tf.reduce_mean(tf.image.ssim(pred, batch['uv_curvature'], max_val=1))
         loss = ce_uv_loss
 
@@ -166,7 +166,7 @@ if __name__ == '__main__':
     parser.add_argument('--name', type=str)
     args = vars(parser.parse_args())
 
-    cce = tf.keras.losses.CategoricalCrossentropy()
+    scce = tf.keras.losses.SparseCategoricalCrossentropy()
 
     resolution_modifier = ''
     if config['resolution'] != 512:

auto_encoder_static.py

@@ -12,13 +12,18 @@ from custom_utils.custom_utils import UV2Mesh, load_model, maskFacesBatch
 import plotly.graph_objects as go
 import plotly.express as ex
 
+if config['is_skirt']:
+    from losses import edge_loss_skirt as edge_loss
+else:
+    from losses import edge_loss
+
 from tf_utils import check_tf_version, allow_memory_growth
-from losses import d_logistic, d_logistic_r1_reg, g_logistic_non_saturating, g_logistic_ns_pathreg, edge_loss, uv_mask_loss
 from dataset import DatasetFactory, StaticDataset
 from custom_utils.custom_utils import batch_to_dict, plot_3d
 from models.cond_encoder import AutoEncoder
 
 
+
 def init_wandb():
     run = wandb.init(project='vae', name=config['name'],
             notes='convnext + convnext, add dense with tanh at the end of reverse convnext,\
@@ -35,7 +40,7 @@ def init_wandb():
     artifact.add_file('convnext/models/convnext_tf.py')
     run.log_artifact(artifact)
 
-uv2mesh = UV2Mesh(res=256)
+uv2mesh = UV2Mesh(res=256, is_skirt=config['is_skirt'])
 
 
 
@@ -93,7 +98,7 @@ def plot_everything(batch: dict, type: str):
     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")
+    F = np.load("custom_utils/smpl_faces" + skirt_modifier + ".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])]
@@ -101,7 +106,7 @@ def plot_everything(batch: dict, type: str):
     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_faces = [np.concatenate((F, 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) }
@@ -162,6 +167,7 @@ if __name__ == '__main__':
     resolution_modifier = ''
     if config['resolution'] != 512:
         resolution_modifier = '_' + str(config['resolution'])
+    skirt_modifier = '_skirt' if config['is_skirt'] else ''
 
     factory = DatasetFactory()
 

configs/config_ae_body.py

@@ -1,8 +1,8 @@
 from configs.config_base import config
 
-config['gpu'] = '1'
+config['gpu'] = '3'
 config['lr'] = 0.0001
-config['name'] = 'ae_body_10_3'
+config['name'] = 'ae_body_10_3_skirt'
 
 config['notes'] = 'convnext + convnext, add dense with tanh at the end of reverse convnext,\
     flattened bottleneck with 128 latent_dim\ remove global orientation\
@@ -13,7 +13,7 @@ config['notes'] = 'convnext + convnext, add dense with tanh at the end of revers
     larger batch size, 32 instead of 4\
     '
 
-config['dataset_in_memory'] = False
+config['dataset_in_memory'] = True
 config['batch_size'] = 32
 
 

configs/config_ae_curv.py

@@ -1,8 +1,9 @@
 from configs.config_base import config
 
-config['gpu'] = '2'
+config['gpu'] = '0'
 config['lr'] = 0.001
-config['name'] = 'x' #'ae_curv_9'
+config['name'] = 'ae_curv_9_skirt' #'ae_curv_9'
+config['notes'] = 'convnext + convnext, 16*8*32, ce loss. flattened latent dim of 128. trained on 100 garments. switched back to sparse CCE, normal one was giving NaN'
 
 config['cond_ae']['is_label'] = True
 config['cond_ae']['upscale_output'] = False

configs/config_ae_static.py

@@ -1,8 +1,8 @@
 from configs.config_base import config
 
-config['gpu'] = '1'
+config['gpu'] = '2'
 config['lr'] = 0.0005
-config['name'] = 'ae_static_7_2'
+config['name'] = 'ae_static_7_2_skirt'
 
 config['cond_ae']['upscale_output'] = False
 

configs/config_base.py

@@ -2,15 +2,17 @@ config = dict()
 
 
 config['batch_size'] = 4
-config['train_dir'] = '/home/workspace/local-data/train_100'
+config['train_dir'] = '/home/workspace/local-data/train_skirt_25_2'
-config['valid_dir'] = '/home/workspace/local-data/train_100'
+config['valid_dir'] = '/home/workspace/local-data/train_skirt_25_2'
 config['use_filled'] = False
 config['resolution'] = 512
+config['is_skirt'] = True
 
 config['latent_dim'] = 128
 
 config['cond_ae'] = dict()
 config['cond_ae']['is_label'] = False
+config['cond_ae']['upscale_output'] = False
 config['cond_ae']['latent_dim'] = config['latent_dim']
 config['cond_ae']['kernel_size'] = 7
 config['cond_ae']['bottleneck_size'] = 64
@@ -18,4 +20,15 @@ config['cond_ae']['encoder_depths'] = [3, 3, 9, 9, 3]
 config['cond_ae']['encoder_dims'] = [32, 64, 128, 256, config['cond_ae']['bottleneck_size']]
 config['cond_ae']['decoder_input'] = (8, 4, config['cond_ae']['bottleneck_size'])
 config['cond_ae']['decoder_depths'] = [3, 3, 9, 3, 3]
-config['cond_ae']['decoder_dims'] = [config['cond_ae']['bottleneck_size'], 256, 128, 64, 16]
+config['cond_ae']['decoder_dims'] = [config['cond_ae']['bottleneck_size'], 256, 128, 64, 16]
+
+config['cond_ae_body'] = dict()
+config['cond_ae_body']['latent_dim'] = config['latent_dim']
+config['cond_ae_body']['kernel_size'] = 7
+config['cond_ae_body']['upscale_output'] = True
+config['cond_ae_body']['bottleneck_size'] = 512
+config['cond_ae_body']['encoder_depths'] = [3, 3, 9, 9, 3, 3]
+config['cond_ae_body']['encoder_dims'] = [32, 64, 128, 256, 256, config['cond_ae_body']['bottleneck_size']]
+config['cond_ae_body']['decoder_input'] = (4, 2, config['cond_ae_body']['bottleneck_size'])
+config['cond_ae_body']['decoder_depths'] = [3, 3, 9, 3]
+config['cond_ae_body']['decoder_dims'] = [config['cond_ae_body']['bottleneck_size'], 128, 64, 16]

configs/config_cvae.py

@@ -1,7 +1,7 @@
 from configs.config_base import config
 
-config['gpu'] = '4'
+config['gpu'] = '2'
-config['name'] = 'cvae_low_4_3_d32'
+config['name'] = 'cvae_19_skirt'
 config['notes'] = '\
 offsets over average unposed garment. remove global orientation (in posed body input and in pose_shape parameters. \
 retrained body encoder with deoriented bodies \
@@ -11,13 +11,13 @@ remove shape from unposed garment\
 training on 100 garments with offsets over the average of these\
 fixed autoencoders!!!!!!!!!\
 cvae only loads encoders instead AEs\
-full res loss\
 '
-config['dataset_in_memory'] = False
+config['dataset_in_memory'] = True
 
 
-config['resolution'] = 32
+config['resolution'] = 512
-
+config['is_skirt'] = True
+config['reg_loss_coeff'] = 0 # set to 0 if don't want to use edge+collision
 
 if config['resolution'] != 512:
   config['notes'] += '\
@@ -50,13 +50,14 @@ if config['resolution'] != 512:
   config['signature_static']['uv_mask' + res_modifier] = (config['resolution'], config['resolution']//2)
 
 if config['resolution'] == 512:
-  config['avg_garment_path'] = 'custom_data/train_100_no_shape_avg_uv.npy'
+  config['avg_garment_path'] = 'custom_data/train_skirt_25_2_no_shape_avg_uv.npy'
 else:
   config['avg_garment_path'] = ''
 
-config['static_model'] = 'checkpoints/ae_static_7_2_x'
+config['static_model'] = 'checkpoints/ae_static_7_2_skirt_x'
-config['body_model'] = 'checkpoints/ae_body_9_x'
+config['body_model'] = 'checkpoints/ae_body_10_3_skirt_x'
-config['curvature_model'] = 'checkpoints/ae_curv_9_x'
+config['curvature_model'] = 'checkpoints/ae_curv_9_skirt_x'
+config['upscale_enhancer'] = 'checkpoints/up_2_3_skirt_x_d32'
 config['freeze_encoders'] = True
 
 config['cvae'] = dict()

convert_to_uv.py

@@ -452,7 +452,7 @@ if __name__ == '__main__':
     parser.add_argument("-s", "--start", help = "start of range", default=0)
     parser.add_argument("-e", "--end", help = "end of range", default=10)
     parser.add_argument("-n", "--name", help = "end of range", default='lhunor')
-    parser.add_argument("-sk", "--skirt", help = "end of range", default=1)
+    parser.add_argument("-sk", "--skirt", help = "end of range", default=0)
     args = parser.parse_args()
 
     # smpl = SMPL('FeatureLines/superSMPL')
@@ -528,9 +528,9 @@ if __name__ == '__main__':
         mask = np.squeeze(sio.loadmat(os.path.join(garment_dir, garment + '_mask.mat'))['mask'].astype(bool))
         # body = readPC2(os.path.join(folder, id, 'Body_posed.pc16'),  True)['V']
         gt = readPC2(os.path.join(folder, id, garment + '_registered.pc16'),  True)['V']
-        curvature = np.load(os.path.join(folder, id, 'curvature_label_' + curvature_label_type + '.npy'))
+        # curvature = np.load(os.path.join(folder, id, 'curvature_label_' + curvature_label_type + '.npy'))
-        mesh_unposed = [np.load(os.path.join(folder, id, 'mesh_unposed',  id + '_' + str(i) + '.npz')) for i in range(gt.shape[0])]
+        # mesh_unposed = [np.load(os.path.join(folder, id, 'mesh_unposed',  id + '_' + str(i) + '.npz')) for i in range(gt.shape[0])]
-        mesh_unposed = np.stack([mesh_unposed[i]['mesh_unposed'] for i in range(gt.shape[0])])
+        # mesh_unposed = np.stack([mesh_unposed[i]['mesh_unposed'] for i in range(gt.shape[0])])
         mesh_unposed_no_shape = [np.load(os.path.join(folder, id, 'mesh_unposed_no_shape',  id + '_' + str(i) + '.npz')) for i in range(gt.shape[0])]
         mesh_unposed_no_shape = np.stack([mesh_unposed_no_shape[i]['mesh_unposed_no_shape'] for i in range(gt.shape[0])])
         # body_no_orient = [np.load(os.path.join(folder, id, 'mesh_posed_body_no_orient',  id + '_' + str(i) + '.npz')) for i in range(gt.shape[0])]
@@ -605,17 +605,17 @@ if __name__ == '__main__':
         # uv_body_static = convert_to_uv(uv_res, m2uv, uv_map, body_static, neighbours=3, mask=None, method='rbf', garment_uv_mask=None).astype(np.float32)
         # uv_body = convert_to_uv_batch(uv_res, m2uv, uv_map, body[:, :, :], neighbours=3, mask=None, method='rbf', garment_uv_mask=None, garment_uv_mask_border=garment_uv_mask_border).astype(np.float32)
         # uv_posed_body_no_orient = convert_to_uv_batch(uv_res, m2uv, uv_map, body_no_orient[:, :, :], neighbours=3, mask=None, method='rbf', garment_uv_mask=None, garment_uv_mask_border=garment_uv_mask_border).astype(np.float32)
-        uv_gt = convert_to_uv_batch(uv_res, m2uv, uv_map, gt[:, :, :], neighbours=3, mask=mask, method='rbf', garment_uv_mask=garment_uv_mask, garment_uv_mask_border=garment_uv_mask_border).astype(np.float32)
+        # uv_gt = convert_to_uv_batch(uv_res, m2uv, uv_map, gt[:, :, :], neighbours=3, mask=mask, method='rbf', garment_uv_mask=garment_uv_mask, garment_uv_mask_border=garment_uv_mask_border).astype(np.float32)
-        uv_curvature = convert_to_uv_batch_label(uv_res, m2uv, uv_map, curvature[:, :], neighbours=3, mask=mask, method='nearest', garment_uv_mask=garment_uv_mask, garment_uv_mask_border=garment_uv_mask_border).astype(np.float32)
+        # uv_curvature = convert_to_uv_batch_label(uv_res, m2uv, uv_map, curvature[:, :], neighbours=3, mask=mask, method='nearest', garment_uv_mask=garment_uv_mask, garment_uv_mask_border=garment_uv_mask_border).astype(np.float32)
-        uv_unposed = convert_to_uv_batch(uv_res, m2uv, uv_map, mesh_unposed[:, :, :], neighbours=3, mask=mask, method='rbf', garment_uv_mask=garment_uv_mask, garment_uv_mask_border=garment_uv_mask_border).astype(np.float32)
+        # uv_unposed = convert_to_uv_batch(uv_res, m2uv, uv_map, mesh_unposed[:, :, :], neighbours=3, mask=mask, method='rbf', garment_uv_mask=garment_uv_mask, garment_uv_mask_border=garment_uv_mask_border).astype(np.float32)
         uv_unposed_no_shape = convert_to_uv_batch(uv_res, m2uv, uv_map, mesh_unposed_no_shape[:, :, :], neighbours=3, mask=mask, method='rbf', garment_uv_mask=garment_uv_mask, garment_uv_mask_border=garment_uv_mask_border).astype(np.float32)
         # uv_shaped_body = convert_to_uv(uv_res, m2uv, uv_map, shaped_body, neighbours=3, mask=None, method='rbf', garment_uv_mask=None).astype(np.float32)
         print(time.perf_counter() - start)
 
         start = time.perf_counter()
-        np.save(os.path.join(folder, id, garment + '_offsets_static_uv.npy'), uv_verts_static)
+        # np.save(os.path.join(folder, id, garment + '_offsets_static_uv.npy'), uv_verts_static)
         # np.save(os.path.join(folder, id, 'body_static_uv.npy'), uv_body_static)
-        np.save(os.path.join(folder, id, 'mask_uv.npy'), garment_uv_mask)
+        # np.save(os.path.join(folder, id, 'mask_uv.npy'), garment_uv_mask)
 
         # os.makedirs(os.path.join(folder, id, 'body_posed_uv'), exist_ok=True)
         # for i in range(uv_body.shape[0]):
@@ -625,17 +625,17 @@ if __name__ == '__main__':
         # for i in range(uv_posed_body_no_orient.shape[0]):
         #     np.savez_compressed(os.path.join(folder, id, 'uv_posed_body_no_orient', garment + '_' + str(i) + '.npz'), uv_posed_body_no_orient=uv_posed_body_no_orient[i, :, :, :])
 
-        os.makedirs(os.path.join(folder, id, 'gt_uv'), exist_ok=True)
+        # os.makedirs(os.path.join(folder, id, 'gt_uv'), exist_ok=True)
-        for i in range(uv_gt.shape[0]):
+        # for i in range(uv_gt.shape[0]):
-            np.savez_compressed(os.path.join(folder, id, 'gt_uv', garment + '_' + str(i) + '.npz'), uv_gt=uv_gt[i, :, :, :])
+        #     np.savez_compressed(os.path.join(folder, id, 'gt_uv', garment + '_' + str(i) + '.npz'), uv_gt=uv_gt[i, :, :, :])
 
-        os.makedirs(os.path.join(folder, id, 'curvature_uv'), exist_ok=True)
+        # os.makedirs(os.path.join(folder, id, 'curvature_uv'), exist_ok=True)
-        for i in range(uv_curvature.shape[0]):
+        # for i in range(uv_curvature.shape[0]):
-            np.savez_compressed(os.path.join(folder, id, 'curvature_uv', garment + '_' + str(i) + '.npz'), curvature_uv=uv_curvature[i, :, :])
+        #     np.savez_compressed(os.path.join(folder, id, 'curvature_uv', garment + '_' + str(i) + '.npz'), curvature_uv=uv_curvature[i, :, :])
 
-        os.makedirs(os.path.join(folder, id, 'uv_unposed'), exist_ok=True)
+        # os.makedirs(os.path.join(folder, id, 'uv_unposed'), exist_ok=True)
-        for i in range(uv_unposed.shape[0]):
+        # for i in range(uv_unposed.shape[0]):
-            np.savez_compressed(os.path.join(folder, id, 'uv_unposed', garment + '_' + str(i) + '.npz'), uv_unposed=uv_unposed[i, :, :, :])
+        #     np.savez_compressed(os.path.join(folder, id, 'uv_unposed', garment + '_' + str(i) + '.npz'), uv_unposed=uv_unposed[i, :, :, :])
 
         os.makedirs(os.path.join(folder, id, 'uv_unposed_no_shape'), exist_ok=True)
         for i in range(uv_unposed_no_shape.shape[0]):

custom_utils/custom_utils.py

@@ -1,7 +1,7 @@
-from faulthandler import disable
 import numpy as np
 import tensorflow as tf
 import plotly.graph_objects as go
+from plotly.subplots import make_subplots
 
 
 def m2uv_array(m2uv):
@@ -19,12 +19,12 @@ def m2uv_array(m2uv):
 
 
 class UV2Mesh():
-    def __init__(self, res=256):
+    def __init__(self, res=256, is_skirt=False):
         super(UV2Mesh, self).__init__()
         self.res = res
-        uv_map = np.load('custom_utils/stretched_uv_map_3.npy')
+        uv_map = np.load('custom_utils/stretched_uv_map_' + ('skirt' if is_skirt else '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 = np.load('custom_utils/m2uv' + ('_skirt' if is_skirt else '') + '.npy', allow_pickle=True).item()
         m2uv = m2uv_array(m2uv)
         self.uv_pixels = self.uv_map_discreet[m2uv]
         # self.batch_size = batch_size
@@ -136,49 +136,56 @@ def maskFacesBatch(V, F, mask):
 
 
 def plot_3d(mesh_gt, mesh_pred_unposed, faces, collision, collision_faces, collision_color):
-    mesh_pred_unposed[:,2] += 1
+    fig = make_subplots(rows=1, cols=3, start_cell="top-left", specs = [ [{ 'type': 'scene' } for _ in range(3)] ])
-    collision[:,2] += 2
-    fig = go.Figure(data=[
-        go.Mesh3d(
-            # 8 vertices of a cube
-            x=mesh_gt[:,0],
-            y=mesh_gt[:,2],
-            z=mesh_gt[:,1],
-            # 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[:,2],
-            z=mesh_pred_unposed[:,1],
-            # 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[:,2],
-            z=collision[:,1],
-            # 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',
+    fig.add_trace(go.Mesh3d(
-                    scene_aspectratio=dict(x=1, y=2, z=1))
+                # 8 vertices of a cube
-    fig.update_layout(
+                x=mesh_gt[:,0],
-        scene = dict(
+                y=mesh_gt[:,2],
-            xaxis = dict(nticks=4, range=[-1,1],),
+                z=mesh_gt[:,1],
-                        yaxis = dict(nticks=4, range=[-1,3],),
+                # i, j and k give the vertices of triangles
-                        zaxis = dict(nticks=4, range=[-1,1],),)
+                i = faces[:,0],
+                j = faces[:,1],
+                k = faces[:,2]
+            ),
+            row=1, col=1)
+    
+    fig.add_trace(go.Mesh3d(
+                # 8 vertices of a cube
+                x=mesh_pred_unposed[:,0],
+                y=mesh_pred_unposed[:,2],
+                z=mesh_pred_unposed[:,1],
+                # i, j and k give the vertices of triangles
+                i = faces[:,0],
+                j = faces[:,1],
+                k = faces[:,2]
+            ),
+            row=1, col=2)
+
+    fig.add_trace(go.Mesh3d(
+                # 8 vertices of a cube
+                x=collision[:,0],
+                y=collision[:,2],
+                z=collision[:,1],
+                # i, j and k give the vertices of triangles
+                i = collision_faces[:,0],
+                j = collision_faces[:,1],
+                k = collision_faces[:,2],
+                vertexcolor=collision_color
+            ),
+            row=1, col=3)
+
+
+    fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=1, z=1),
+                    scene2_aspectmode='manual', scene2_aspectratio=dict(x=1, y=1, z=1),
+                    scene3_aspectmode='manual', scene3_aspectratio=dict(x=1, y=1, z=1),
+                    )
+    helper = dict(xaxis = dict(nticks=4, range=[-1,1],),
+                yaxis = dict(nticks=4, range=[-1,1],),
+                zaxis = dict(nticks=4, range=[-1,1],),)
+    fig.update_layout(scene = helper,
+                        scene2 = helper,
+                        scene3 = helper,
                         )
 
     return fig

cvae.py

@@ -26,10 +26,15 @@ import plotly.graph_objects as go
 import plotly.express as ex
 from PIL import Image
 
+if config['is_skirt']:
+    from losses import edge_loss_skirt as edge_loss
+else:
+    from losses import edge_loss
+
 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 edge_loss, uv_mask_loss, bend_loss, kl_loss
+from losses import uv_mask_loss, bend_loss, kl_loss
 # from losses import collision_penalty as collision_penalty
 from losses import collision_loss as collision_penalty
 from custom_utils.custom_utils import batch_to_dict, load_model, plot_3d
@@ -52,7 +57,7 @@ def init_wandb():
     artifact.add_file('convnext/models/convnext_tf.py')
     run.log_artifact(artifact)
 
-uv2mesh = UV2Mesh(res=256)
+uv2mesh = UV2Mesh(res=256, is_skirt=config['is_skirt'])
 
 
 def log_normal_pdf(sample, mean, logvar, raxis=1):
@@ -72,10 +77,7 @@ def kl_loss(mean, logvar):
 def train_step(batch):
 
     with tf.GradientTape() as tape:
-        uv_curvature = tf.one_hot(batch['uv_curvature'], 7)
+        pred, mean, logvar = model([batch['uv_unposed'], batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']], training=True)
-
-        pred, mean, logvar, z = model([batch['uv_unposed'], batch['uv_static'], batch['uv_body_posed'], uv_curvature, batch['uv_mask']],
-                                      training=True, use_mask=True)
 
         mesh_pred = uv2mesh(pred)
         mesh_orig = batch['mesh_unposed']
@@ -94,7 +96,7 @@ def train_step(batch):
         ms_ssim = tf.image.ssim_multiscale(pred + 1, batch['uv_unposed'] + 1, max_val=3, power_factors=(0.5, 0.2, 0.1, 0.1, 0.1))
         # 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))
 
-        reconstr_loss = l2_loss + e_loss + c_loss
+        reconstr_loss = l2_loss + (e_loss + c_loss) * config['reg_loss_coeff']
         reconstr_loss = tf.reduce_mean(reconstr_loss)
         kl = kl_loss(mean, logvar)
         loss = reconstr_loss + balancing_coeff * kl
@@ -152,8 +154,7 @@ def valid_step(batch):
 
 
 def plot_everything(batch: dict, type: str):
-    uv_curvature = tf.one_hot(batch['uv_curvature'], 7)
+    pred, mean, logvar = model([batch['uv_unposed'], batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']], training=False)
-    pred, mean, logvar, z = model([batch['uv_unposed'], batch['uv_static'], batch['uv_body_posed'], uv_curvature, batch['uv_mask']], training=False, use_mask=True)
     
     mesh_pred = uv2mesh(pred).numpy()
     mesh_orig = batch['mesh_unposed']
@@ -169,7 +170,7 @@ def plot_everything(batch: dict, type: str):
     gender = tf.cast(tf.greater_equal(shape[:, 10] , 0.5), tf.int32)
     shape_offsets = tf.stack([tf.gather(tf.tensordot(shape[i,:10][tf.newaxis, :], smpl.shapedirs[gender[i]], axes=[[1], [2]])[0,:,:], smpl.to_keep, axis=0) for i in range(shape.shape[0])])
 
-    F = np.load("custom_utils/smpl_faces.npy")
+    F = np.load("custom_utils/smpl_faces" + skirt_modifier + ".npy")
     mask = batch['mesh_mask'].numpy().astype(bool)
     garment_unposed, faces = maskFacesBatch(mesh_pred, F, mask)
     garment_unposed = [garment_unposed[i] + shape_offsets.numpy()[i, mask[i], :] for i in range(mask.shape[0])]
@@ -198,7 +199,7 @@ def sample(folder):
     load_model(model,folder)
     count = 0
     for batch in dataset_train:
-        batch = static_dataset_train.expand_batch(batch_to_dict(batch, signature_dynamic))
+        batch = static_dataset_train.expand_batch(batch_to_dict(batch, config['signature_dynamic']))
         uv_curvature = tf.one_hot(batch['uv_curvature'], 7)
         nr_of_samples=2
 
@@ -322,6 +323,7 @@ if __name__ == '__main__':
     args = vars(parser.parse_args())
 
     smpl = SMPL()
+    skirt_modifier = '_skirt' if config['is_skirt'] else ''
 
     factory = DatasetFactory()
 

cvae_low.py

@@ -27,7 +27,12 @@ import plotly.express as ex
 from PIL import Image
 
 from tf_utils import check_tf_version, allow_memory_growth, split_gpu_for_testing
-from losses import d_logistic, d_logistic_r1_reg, g_logistic_non_saturating, g_logistic_ns_pathreg, edge_loss, uv_mask_loss, bend_loss
+
+if config['is_skirt']:
+    from losses import edge_loss_skirt as edge_loss
+else:
+    from losses import edge_loss
+
 # from losses import collision_penalty as collision_penalty
 from losses import collision_loss as collision_penalty
 from custom_utils.custom_utils import batch_to_dict, load_model, plot_3d
@@ -51,7 +56,7 @@ def init_wandb():
     artifact.add_file('convnext/models/convnext_tf.py')
     run.log_artifact(artifact)
 
-uv2mesh = UV2Mesh(res=256)
+uv2mesh = UV2Mesh(res=256, is_skirt=config['is_skirt'])
 
 
 def log_normal_pdf(sample, mean, logvar, raxis=1):
@@ -71,10 +76,8 @@ def kl_loss(mean, logvar):
 def train_step(batch):
 
     with tf.GradientTape() as tape:
-        uv_curvature = tf.one_hot(batch['uv_curvature'], 7)
+        pred, mean, logvar = model([batch['uv_unposed' + res_modifier], batch['uv_static' + res_modifier], batch['uv_body_posed' + res_modifier], batch['uv_curvature' + res_modifier],
-        uv_curvature_low = tf.one_hot(batch['uv_curvature' + res_modifier], 7)
+                                       batch['uv_mask' + res_modifier], batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']])
-        pred, mean, logvar, z = model([batch['uv_unposed' + res_modifier], batch['uv_static' + res_modifier], batch['uv_body_posed' + res_modifier], uv_curvature_low,
-                                       batch['uv_mask' + res_modifier], batch['uv_static'], batch['uv_body_posed'], uv_curvature, batch['uv_mask']], use_mask=True)
 
         pred_enhanced = enhance(upscale_enhancer, pred, batch['uv_mask' + res_modifier], training=False)
         pred_full_res = tf.image.resize(pred_enhanced, size=(512, 256), method=tf.image.ResizeMethod.LANCZOS5) * batch['uv_mask'][:,:,:,tf.newaxis]
@@ -91,22 +94,18 @@ def train_step(batch):
         c_loss = collision_penalty(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask'])
         l1_loss = tf.reduce_sum(tf.reduce_sum((tf.math.abs(batch['uv_unposed' + res_modifier] - pred)), axis=-1), axis=[1, 2])
         l2_loss = 1e3 * tf.reduce_mean(tf.reduce_sum(((batch['uv_unposed' + res_modifier] - pred) ** 2), axis=-1), axis=[1, 2])
-        l2_loss_full = 1e3 * tf.reduce_mean(tf.reduce_sum(((batch['uv_unposed'] - pred_full_res) ** 2), axis=-1), axis=[1, 2])
         if resolution == 32:
             ms_ssim = tf.image.ssim(pred + 1, batch['uv_unposed' + res_modifier] + 1, max_val=3) 
-            ms_ssim_full = tf.image.ssim(pred_full_res + 1, batch['uv_unposed'] + 1, max_val=3) 
         else:
             ms_ssim = tf.image.ssim_multiscale(pred + 1, batch['uv_unposed' + res_modifier] + 1, max_val=3, power_factors=(0.5, 0.3, 0.1, 0.1)) 
 
-        reconstr_loss = l2_loss_full #+ e_loss + c_loss
+        reconstr_loss = l2_loss #+ e_loss + c_loss
         reconstr_loss = tf.reduce_mean(reconstr_loss)
         kl = kl_loss(mean, logvar)
         loss = reconstr_loss + balancing_coeff * kl
         ms_ssim = tf.reduce_mean(ms_ssim)
-        ms_ssim_full = tf.reduce_mean(ms_ssim_full)
         l1_loss = tf.reduce_mean(l1_loss)
         l2_loss = tf.reduce_mean(l2_loss)
-        l2_loss_full = tf.reduce_mean(l2_loss_full)
         e_loss = tf.reduce_mean(e_loss)
         c_loss = tf.reduce_mean(c_loss)
         mean = tf.reduce_mean(mean)
@@ -116,7 +115,6 @@ def train_step(batch):
     gradients = tape.gradient(loss, model.trainable_variables)
     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/l2_loss_full_res': l2_loss_full, 'train/ms_ssim_full_res': ms_ssim_full,
              'train/reconstr_loss': reconstr_loss,
              'train/mean': mean, 'train/logvar': logvar, 'train/kl_loss': kl, 'train/var': np.e ** logvar,
              'train/reg_coeff': reg_coeff, 'train/l2_loss': l2_loss }
@@ -163,29 +161,28 @@ def valid_step(batch):
 
 
 def plot_everything(batch: dict, type: str):
-    uv_curvature = tf.one_hot(batch['uv_curvature'], 7)
+    pred, mean, logvar = model([batch['uv_unposed' + res_modifier], batch['uv_static' + res_modifier], batch['uv_body_posed' + res_modifier], batch['uv_curvature' + res_modifier],
-    uv_curvature_low = tf.one_hot(batch['uv_curvature' + res_modifier], 7)
+                                       batch['uv_mask' + res_modifier], batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']])
-    pred, mean, logvar, z = model([batch['uv_unposed' + res_modifier], batch['uv_static' + res_modifier], batch['uv_body_posed' + res_modifier], uv_curvature_low,
-                                    batch['uv_mask' + res_modifier], batch['uv_static'], batch['uv_body_posed'], uv_curvature, batch['uv_mask']], use_mask=True)
     
     pred = enhance(upscale_enhancer, pred, batch['uv_mask' + res_modifier], training=False)
-    mesh_pred = uv2mesh(pred).numpy()
+    pred_full_res = tf.image.resize(pred, size=(512, 256), method=tf.image.ResizeMethod.LANCZOS5) * batch['uv_mask'][:,:,:,tf.newaxis]
+    mesh_pred = uv2mesh(pred_full_res).numpy()
     mesh_orig = batch['mesh_unposed']
-    
+
-    orig_img = batch['uv_unposed' + res_modifier]
+    orig_img = (batch['uv_unposed'].numpy() + 1) / 2
-    uv_static_img = batch['uv_static' + res_modifier]
+    uv_static_img = (batch['uv_static'].numpy() + 1) / 2
-    uv_body_posed_img = batch['uv_body_posed' + res_modifier]
+    uv_body_posed_img = (batch['uv_body_posed'].numpy() + 1) / 2
-    pred_img = pred.numpy()
+    uv_curvature = (np.tile(batch['uv_curvature'].numpy().argmax(-1)[:,:,:,np.newaxis], [1,1,1,3]) / 6)
-    pred_img[pred_img > 1] = 0
+    pred_img = pred_full_res.numpy()
-    pred_img[pred_img < -1] = 0
+    pred_img = (pred_img - pred_img.min(axis=(1,2,3))[:,np.newaxis,np.newaxis,np.newaxis]) / (pred_img.max(axis=(1,2,3))[:,np.newaxis,np.newaxis,np.newaxis] - pred_img.min(axis=(1,2,3))[:,np.newaxis,np.newaxis,np.newaxis])
-    img = tf.concat((uv_static_img, uv_body_posed_img, orig_img, pred_img), axis=2)
+    img = tf.concat((uv_static_img, uv_body_posed_img, uv_curvature, orig_img, pred_img), axis=2)
-    imgs = [wandb.Image((img[i] + 1) / 2.0) for i in range(4)]
+    imgs = [wandb.Image(img[i]) for i in range(4)]
-    
+
     shape = batch['pose_shape'][:, 72:]
     gender = tf.cast(tf.greater_equal(shape[:, 10] , 0.5), tf.int32)
     shape_offsets = tf.stack([tf.gather(tf.tensordot(shape[i,:10][tf.newaxis, :], smpl.shapedirs[gender[i]], axes=[[1], [2]])[0,:,:], smpl.to_keep, axis=0) for i in range(shape.shape[0])])
 
-    F = np.load("custom_utils/smpl_faces.npy")
+    F = np.load("custom_utils/smpl_faces" + skirt_modifier + ".npy")
     mask = batch['mesh_mask'].numpy().astype(bool)
     garment_unposed, faces = maskFacesBatch(mesh_pred, F, mask)
     garment_unposed = [garment_unposed[i] + shape_offsets.numpy()[i, mask[i], :] for i in range(mask.shape[0])]
@@ -366,11 +363,13 @@ if __name__ == '__main__':
     
     smpl = SMPL()
 
-    resolution = 32
+    resolution = config['resolution']
     res_modifier = ''
     if resolution != 512:
         res_modifier = '_' + str(resolution)
 
+    skirt_modifier = '_skirt' if config['is_skirt'] else ''
+
     factory = DatasetFactory()
 
     dataset_train = factory.get_dataset(folder=config['train_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'], use_filled=config['use_filled'], resolutions=[512, config['resolution']], in_memory=config['dataset_in_memory'])
@@ -380,7 +379,7 @@ if __name__ == '__main__':
 
     model = CVAE(config)
 
-    upscale_enhancer = get_upscale_enhancer('checkpoints/up_2_3', res=resolution)
+    upscale_enhancer = get_upscale_enhancer(config['upscale_enhancer'], res=resolution)
 
     optimizer = tf.optimizers.Adam(learning_rate=config['lr'])
     balancing_coeff = config['starting_balancing_coeff']

losses.py

@@ -142,6 +142,10 @@ def collision_loss(V, B, mask, thr=.004):
 laplacian_loss = sio.loadmat("custom_utils/ssmpl0.mat")['L'].tocoo()
 edges = np.stack(np.nonzero(laplacian_loss.todense())).transpose()
 
+laplacian_loss_skirt = sio.loadmat("custom_utils/ssmpl_skirt0.mat")['L'].tocoo()
+edges_skirt = np.stack(np.nonzero(laplacian_loss.todense())).transpose()
+
+
 def get_edge_mask(masks):
     return tf.reduce_min(tf.gather(masks, edges, axis=1), axis=-1)
 
@@ -158,6 +162,22 @@ def edge_loss(x, y, mask):
     loss = tf.reduce_sum(d_e ** 2, axis=-1)# / tf.reduce_sum(e_mask, axis=-1)
     return loss
 
+def get_edge_mask_skirt(masks):
+    return tf.reduce_min(tf.gather(masks, edges_skirt, axis=1), axis=-1)
+
+def edge_loss_skirt(x, y, mask):
+    e_mask = get_edge_mask_skirt(mask)
+    x_e = tf.gather(x, edges_skirt[:,0], axis=1) - tf.gather(x, edges_skirt[:,1], axis=1)
+    x_e = tf.reduce_sum(x_e ** 2, -1)
+    x_e = tf.sqrt(x_e + 1e-9)
+    y_e = tf.gather(y, edges_skirt[:,0], axis=1) - tf.gather(y, edges_skirt[:,1], axis=1)
+    y_e = tf.reduce_sum(y_e ** 2, -1)
+    y_e = tf.sqrt(y_e + 1e-9)
+    d_e = y_e - x_e
+    d_e = d_e * e_mask
+    loss = tf.reduce_sum(d_e ** 2, axis=-1)# / tf.reduce_sum(e_mask, axis=-1)
+    return loss
+
 
 class PairwiseDistance(tf.keras.layers.Layer):
     def __init__(self, **kwargs):

models/cvae_model.py

@@ -22,7 +22,7 @@ class CVAE(tf.keras.Model):
         if config['static_model'] != '':
             load_model(self.static_enc, config['static_model'] + '_enc')
 
-        self.body_enc : AEEncoder = AEEncoder(config['cond_ae']['latent_dim'], config['cond_ae']['encoder_depths'], config['cond_ae']['encoder_dims'], config['cond_ae']['bottleneck_size'], config['cond_ae']['kernel_size'])
+        self.body_enc : AEEncoder = AEEncoder(config['cond_ae_body']['latent_dim'], config['cond_ae_body']['encoder_depths'], config['cond_ae_body']['encoder_dims'], config['cond_ae_body']['bottleneck_size'], config['cond_ae_body']['kernel_size'])
         if config['body_model'] != '':
             load_model(self.body_enc, config['body_model'] + '_enc')
 
@@ -49,38 +49,6 @@ class CVAE(tf.keras.Model):
             self.avg_garment = None
 
 
-
-    # @tf.function
-    def sample_single(self, uv_template, uv_body_posed, uv_curv, uv_mask, nr_of_samples=4, fix_noise=False):
-        if not fix_noise:
-            eps = tf.random.normal(shape=(nr_of_samples, self.latent_dim))
-        else:
-            eps = tf.random.normal(shape=(1, self.latent_dim))
-            eps = tf.tile(eps, [nr_of_samples,1])
-
-        uv_template = tf.tile(tf.transpose(uv_template, [1,2,0])[tf.newaxis,:,:,:], [nr_of_samples,1,1,1])
-        uv_body_posed = tf.tile(tf.transpose(uv_body_posed, [1,2,0])[tf.newaxis,:,:,:], [nr_of_samples,1,1,1])
-        if len(uv_curv.shape) == 3:
-            uv_curv = tf.tile(uv_curv[tf.newaxis,:,:,:], [nr_of_samples,1,1,1])
-        uv_mask = tf.tile(uv_mask[tf.newaxis,:,:], [nr_of_samples,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_cond = self.cond_encoder(z_cond)
-
-        z = tf.concat([eps, z_cond], axis=-1)
-
-        x = self.decoder(z)
-        x = x + self.avg_garment[tf.newaxis,:,:,:]
-        x = x * uv_mask[:,:,:,tf.newaxis]
-
-        return x
-
-
     def sample(self, uv_template, uv_body_posed, uv_curv, uv_mask=None, noise=None):
         if noise is None:
             eps = tf.random.normal(shape=(uv_template.shape[0], self.latent_dim))
@@ -110,24 +78,9 @@ class CVAE(tf.keras.Model):
         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 encode_conds(self, uv_template, uv_body_posed, uv_curv):
-        if uv_template.shape[1] == 3:
-            uv_template = tf.transpose(uv_template, [0,2,3,1])
-        if uv_body_posed.shape[1] == 3:
-            uv_body_posed = tf.transpose(uv_body_posed, [0,2,3,1])
-        if uv_curv.shape[1] == 3:
-            uv_curv = tf.transpose(uv_curv, [0,2,3,1])
-
-        z_static = self.static_enc(uv_template)
         z_body = self.body_enc(uv_body_posed)
+        z_static = self.static_enc(uv_template)
         z_curv = self.curv_enc(uv_curv)
 
         z_cond = tf.concat([z_static, z_body, z_curv], axis=-1)
@@ -136,20 +89,14 @@ class CVAE(tf.keras.Model):
         return z_cond
 
 
-    def call(self, inputs, training=None, use_mask=True):
+    def call(self, inputs, training=None):
 
-        if use_mask:
+        if self.size == 512:
-            if self.size == 512:
+            uv_unposed, uv_template, uv_body_posed, uv_curv, uv_mask = inputs
-                uv_unposed, uv_template, uv_body_posed, uv_curv, uv_mask = inputs
+            uv_template_high, uv_body_posed_high, uv_curv_high = uv_template, uv_body_posed, uv_curv
-                uv_template_high, uv_body_posed_high, uv_curv_high = uv_template, uv_body_posed, uv_curv
-            else:
-                uv_unposed, uv_template, uv_body_posed, uv_curv, uv_mask, uv_template_high, uv_body_posed_high, uv_curv_high, uv_mask_high = inputs
         else:
-            if self.size == 512:
+            uv_unposed, uv_template, uv_body_posed, uv_curv, uv_mask, uv_template_high, uv_body_posed_high, uv_curv_high, uv_mask_high = inputs
-                uv_unposed, uv_template, uv_body_posed, uv_curv = inputs
+
-                uv_template_high, uv_body_posed_high, uv_curv_high = uv_template, uv_body_posed, uv_curv
-            else:
-                uv_unposed, uv_template, uv_body_posed, uv_curv, uv_template_high, uv_body_posed_high, uv_curv_high = inputs
 
         x = tf.concat([uv_unposed, uv_template, uv_body_posed, uv_curv], axis=-1)
         mean, logvar = self.encoder(x, training=training)
@@ -161,8 +108,8 @@ class CVAE(tf.keras.Model):
 
         if self.avg_garment is not None:
             x = x + self.avg_garment[tf.newaxis,:,:,:]
-        if use_mask:
+ 
-            x = x * uv_mask[:,:,:,tf.newaxis]
+        x = x * uv_mask[:,:,:,tf.newaxis]
 
-        return x, mean, logvar, z
+        return x, mean, logvar
 

models/upscaler.py

@@ -75,9 +75,6 @@ class UpscalerBasic(tf.keras.Model):
     def call(self, inputs, training=None):
         img, mask = inputs
 
-        if img.shape[1] == 3:
-            img = tf.transpose(img, [0,2,3,1])
-
         x = self.flatten(img)
         x = self.seq(x, training=training)
         x = self.reshape(x)

unpose_skirt.py

@@ -0,0 +1,391 @@
+import os
+import sys
+import time
+import argparse
+
+import numpy as np
+import scipy.io as sio
+import scipy as sp
+import plotly.graph_objects as go
+from scipy.spatial.transform import Rotation
+
+from tqdm import tqdm
+
+
+os.environ["CUDA_VISIBLE_DEVICES"] = "1"
+import tensorflow as tf
+gpus = tf.config.experimental.list_physical_devices('GPU')
+if gpus:
+    try:
+        for gpu in gpus:
+            tf.config.experimental.set_memory_growth(gpu, True)
+
+    except RuntimeError as e:
+        print(e)
+
+
+from custom_utils.smpl import SMPL
+from custom_utils.smpl_2 import SMPL as SMPL_2
+from custom_utils.supersmpl_io import readPC2
+
+
+
+# run = wandb.init(config=tf.compat.v1.flags.FLAGS, sync_tensorboard=True, project='unpose', name='new_smpl_2', notes='')
+# artifact = wandb.Artifact('code', type='code')
+# artifact.add_file('Data/unpose.py')
+# artifact.add_file('Data/verify_unposed.py')
+# artifact.add_file('FeatureLines/superSMPL/smpl.py')
+# artifact.add_file('FeatureLines/superSMPL/smpl_2.py')
+# run.log_artifact(artifact)
+
+
+def read_garment(folder, id):
+    garment_dir = os.path.join(folder, id)
+
+    garments = [f.replace('.pc16','') for f in os.listdir(garment_dir) if f.find('_') == -1 and f != 'info.mat' and not f.endswith('.obj')]
+    assert len (garments) == 1
+        
+    garment = garments[0]
+
+    posed = readPC2(os.path.join(folder, id, garment + '_registered.pc16'), True)['V']
+    info = sio.loadmat(os.path.join(folder, id, 'info.mat'))
+    pose, shape, gender = info['poses'].astype(np.float32), info['shape'], info['gender']
+    shape = np.squeeze(np.concatenate((shape, gender), axis=1)).astype(np.float32)
+    mask = np.load(os.path.join(folder, id, garment + '_mask.npy'))
+    offsets = sio.loadmat(os.path.join(folder, id, garment + '_offsets_static.mat'))['label'][:,:3]
+    template = sio.loadmat(os.path.join(folder, id, 'info_body.mat'))['B']
+    mesh_static = (template + offsets) * mask[:, np.newaxis]
+
+    posed, mesh_static = posed[:, mask, :], mesh_static[mask, :]
+    mask_idx = np.array(range(14475))[mask]
+
+    return posed, mask, pose, shape, mesh_static, mask_idx
+
+
+def plot_3d(mesh_orig, mesh_pred, unposed, faces):
+    _mesh_pred = mesh_pred.copy()
+    _unposed = unposed.copy()
+    _mesh_pred[:,1] += 1
+    _unposed[:,1] += 3
+    fig = go.Figure(data=[
+        go.Mesh3d(
+            # 8 vertices of a cube
+            x=mesh_orig[:,0],
+            y=mesh_orig[:,1],
+            z=mesh_orig[:,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],
+        ),
+        go.Mesh3d(
+            # 8 vertices of a cube
+            x=_unposed[:,0],
+            y=_unposed[:,1],
+            z=_unposed[:,2],
+            # i, j and k give the vertices of triangles
+            i = faces[:,0],
+            j = faces[:,1],
+            k = 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
+
+
+
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument("-f", "--folder", help = "end of range", default='/home/workspace/data/dataset/train_skirt')
+    parser.add_argument("-s", "--start", help = "start of range", default=0)
+    parser.add_argument("-e", "--end", help = "end of range", default=1)
+    parser.add_argument("-n", "--name", help = "end of range", default='lhunor')
+    args = parser.parse_args()
+
+    F = np.load('custom_utils/smpl_faces_skirt.npy')
+    folder = args.folder
+    # smpl = SMPL()
+    smpl_2 = SMPL_2('custom_utils')
+    iterations = 10000
+
+    ids = os.listdir(folder)[int(args.start):int(args.end)]
+
+    for id in tqdm(ids):
+        if not os.path.isdir(os.path.join(folder, id)):
+            continue
+
+        # if len(id) == 5:
+        #     continue
+
+        os.makedirs(os.path.join(folder, id, 'mesh_unposed'), exist_ok=True)
+        
+        posed, mask, pose, shape, mesh_static, mask_idx = read_garment(folder, id)
+        gender = int(shape[-1] > 0)
+        init =  Rotation.from_euler('zyz', [90,90,90], degrees=True).apply(mesh_static.copy()).astype(np.float32)
+
+        existing_frames = os.listdir(os.path.join(folder, id, 'mesh_unposed'))
+        if len(existing_frames) == posed.shape[0]:
+            continue
+
+
+        R = smpl_2.run_rodrigues(np.transpose(pose, [1,0]))
+
+        frames = posed.shape[0]
+
+        # rot = Rotation.from_euler('z', [90], degrees=True)
+        # rest_pose = np.zeros((1,72), dtype=np.float32)
+        # R_rest = smpl.run_rodrigues(rest_pose)
+        # body = smpl.run(R_rest, shape[tf.newaxis,:])[0].numpy()[:, mask,:]
+
+        # # verts = smpl.run(R, tf.tile(shape[tf.newaxis,:], [frames,1]))[0].numpy()[:,mask,:]
+        # # verts = tf.stack([rot.apply(verts[i]).astype(np.float32) for i in range(frames)])
+        # verts = tf.tile(body, [frames,1,1])
+
+        # start = time.perf_counter()
+        # result = smpl.run_garment(R, tf.tile(shape[tf.newaxis,:], [frames,1]), verts, filter_idx=mask_idx)
+        # orig_time = time.perf_counter() - start
+        # print(orig_time)
+
+        # start = time.perf_counter()
+        # result_2 = smpl_2.run_garment(R, shape[:10], gender, verts, filter_idx=mask_idx)
+        # new_time = time.perf_counter() - start
+        # print(new_time)
+
+        # # test = np.abs(result.numpy() - result_2.numpy()) > 1e-4
+        # # x = np.sum(test, axis=-1).sum(-1).sum(-1) 
+        # euclid_error = tf.reduce_mean(tf.math.sqrt(tf.reduce_sum((result - result_2) ** 2, axis=-1))).numpy()
+
+        # print('speedup: ' + str(orig_time/new_time) + ' euclid_error: ' + str(euclid_error))
+
+        # # continue
+
+        # batch = np.random.choice(np.array(range(frames)))
+        # _, faces = _maskFaces(np.zeros((14475,3)), F, mask)
+        # fig = go.Figure(data=[
+        #     go.Mesh3d(
+        #         # 8 vertices of a cube
+        #         x=result[batch,:,0],
+        #         y=result[batch,:,1],
+        #         z=result[batch,:,2],
+        #         # i, j and k give the vertices of triangles
+        #         i = faces[:,0],
+        #         j = faces[:,1],
+        #         k = faces[:,2],
+        #         text='speedup: ' + str(orig_time/new_time) + ' euclid_error: ' + str(euclid_error),
+        #     ),
+        #     go.Mesh3d(
+        #         # 8 vertices of a cube
+        #         x=result_2[batch,:,0],
+        #         y=result_2[batch,:,1],
+        #         z=result_2[batch,:,2],
+        #         # i, j and k give the vertices of triangles
+        #         i = faces[:,0],
+        #         j = faces[:,1],
+        #         k = faces[:,2],
+        #         text='speedup: ' + str(orig_time/new_time) + ' euclid_error: ' + str(euclid_error),
+        #     )],
+        #     )
+
+        # wandb.log({id: fig})
+
+        # continue
+
+
+
+        lr_schedule = tf.keras.optimizers.schedules.ExponentialDecay(1000.0, decay_steps=40000, decay_rate=0.1, staircase=False)
+        optimizer = tf.keras.optimizers.SGD(learning_rate=lr_schedule)
+        unposed = tf.Variable(tf.tile(init[tf.newaxis,:,:], [frames,1,1]))
+
+
+        @tf.function
+        def skinning(verts):
+            result = smpl_2.run_garment(R, shape[:10], gender, verts, filter_idx=mask_idx)
+            return result
+
+
+        @tf.function
+        def error():
+            pred = skinning(unposed)
+            loss = tf.reduce_mean(tf.math.abs(pred - posed))
+
+            return loss
+
+        start = time.perf_counter()
+        euclid_error = 10000
+
+    
+        while True:
+            optimizer.minimize(error, [unposed])
+
+            i = optimizer.iterations.numpy()
+            if i % 1000 == 0:
+                pred = skinning(unposed)
+                # loss = tf.reduce_mean(tf.reduce_sum(tf.math.abs(pred - posed) ** 2, axis=-1))
+                loss = tf.reduce_mean(tf.math.abs(pred - posed))
+                euclid_error = tf.reduce_mean(tf.math.sqrt(tf.reduce_sum((posed - pred) ** 2, axis=-1)))
+                print('id: ' + id  + ' iteration: ' + str(i) + ' loss: ' + str(loss.numpy()) + ' error: ' + str(euclid_error.numpy()) + ' time: ' + str(time.perf_counter() - start))
+                start = time.perf_counter()
+            
+            if euclid_error < 0.0001:
+                break
+
+            if euclid_error < 0.0005 and i > 15000:
+                break
+
+            if euclid_error < 0.001 and i > 30000:
+                break
+
+            if i > 40000:
+                break
+
+        # b = np.random.choice(np.array(range(frames)))
+        # _, faces = _maskFaces(np.zeros((14475,3)), F, mask)
+        # plot_3d(posed[b], pred.numpy()[b], unposed.numpy()[b], faces).show()
+
+        result = np.zeros((frames, 14475, 3), dtype = np.float32)
+        result[:, mask,:] = unposed.numpy()
+        for frame in range(frames):
+            np.savez_compressed(os.path.join(folder, id, 'mesh_unposed', id + '_' + str(frame) + '.npz'), mesh_unposed=result[frame])
+        
+        # batch = 0
+        # go.Figure(data=[
+        #     go.Mesh3d(
+        #         # 8 vertices of a cube
+        #         x=posed[batch,:,0],
+        #         y=posed[batch,:,1],
+        #         z=posed[batch,:,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=pred[batch,:,0],
+        #         y=pred[batch,:,1],
+        #         z=pred[batch,:,2],
+        #         # i, j and k give the vertices of triangles
+        #         i = faces[:,0],
+        #         j = faces[:,1],
+        #         k = faces[:,2],
+        #     )],
+        #     ).show()
+
+
+        continue
+
+
+
+        offset = 1
+
+
+        for frame in range(posed.shape[0]):
+   
+            lr_schedule = tf.keras.optimizers.schedules.ExponentialDecay(10.0, decay_steps=10000, decay_rate=0.1, staircase=False)
+            optimizer = tf.keras.optimizers.SGD(learning_rate=lr_schedule)
+            unposed = tf.Variable(np.random.rand(offset, posed.shape[1], 3).astype(np.float32) * 2 - 1)
+            # unposed = tf.Variable(init)
+
+            @tf.function
+            def skinning(verts):
+                # verts = smpl.run(R[frame:frame+offset], tf.tile(shape[tf.newaxis,:], [offset,1]))[0].numpy()[:,mask,:]
+                
+
+                # start = time.perf_counter()
+                result = smpl.run_garment(R[frame:frame+offset], tf.tile(shape[tf.newaxis,:], [offset,1]), verts, filter_idx=mask_idx)
+                # print(time.perf_counter() - start)
+
+                # start = time.perf_counter()
+                # result_2 = smpl_2.run_garment(R[frame:frame+offset], shape[:10], gender, verts, filter_idx=mask_idx)
+                # print(time.perf_counter() - start)
+
+                # test = np.abs(result.numpy() - result_2.numpy()) > 1e-4
+                # x = np.sum(test, axis=-1).sum(-1).sum(-1) 
+                # euclid_error = tf.reduce_mean(tf.math.sqrt(tf.reduce_sum((result - result_2) ** 2, axis=-1)))
+
+                # batch = 25
+                # _, faces = _maskFaces(np.zeros((14475,3)), F, mask)
+                # go.Figure(data=[
+                #     go.Mesh3d(
+                #         # 8 vertices of a cube
+                #         x=result[batch,:,0],
+                #         y=result[batch,:,1],
+                #         z=result[batch,:,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=result_2[batch,:,0],
+                #         y=result_2[batch,:,1],
+                #         z=result_2[batch,:,2],
+                #         # i, j and k give the vertices of triangles
+                #         i = faces[:,0],
+                #         j = faces[:,1],
+                #         k = faces[:,2],
+                #     )]).show()
+
+                return result
+
+            @tf.function
+            def error():
+                pred = skinning(unposed)
+                loss = tf.reduce_mean(tf.math.abs(pred - posed[frame]))
+
+                return loss
+
+            start = time.perf_counter()
+            euclid_error = 10000
+
+            for i in range(iterations):
+                optimizer.minimize(error, [unposed])
+                if i % 50 == 0:
+                    pred = skinning(unposed)
+                    # loss = tf.reduce_mean(tf.reduce_sum(tf.math.abs(pred - posed) ** 2, axis=-1))
+                    loss = tf.reduce_mean(tf.math.abs(pred - posed[frame]))
+                    euclid_error = tf.reduce_mean(tf.math.sqrt(tf.reduce_sum((posed[frame] - pred) ** 2, axis=-1)))
+                    print('id: ' + id + ' frame: ' + str(frame) + ' iteration: ' + str(i) + ' loss: ' + str(loss.numpy()) + ' error: ' + str(euclid_error.numpy()) + ' time: ' + str(time.perf_counter() - start))
+                    start = time.perf_counter()
+                
+                if euclid_error < 0.001:
+                    break
+
+            result = np.zeros((14475, 3), dtype = np.float32)
+            result[mask,:] = unposed.numpy()
+            np.savez_compressed(os.path.join(folder, id, 'mesh_unposed', id + '_' + str(frame) + '.npz'), mesh_unposed=result)
+            
+            init = unposed
+
+            # pred, unposed =  pred.numpy(), unposed.numpy()
+            # _, faces = _maskFaces(np.zeros((14475,3)), F, mask)
+            # k = 0
+            # plot_3d(posed[k], pred[k], unposed[k], faces).show()
+
+        # print()
+
+
+    

upscale_test.py

@@ -1,5 +1,18 @@
 import os
-os.environ["CUDA_VISIBLE_DEVICES"] = "2"
+os.environ["CUDA_VISIBLE_DEVICES"] = "4"
+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
@@ -7,8 +20,7 @@ import wandb
 import random
 import cv2
 import numpy as np
-import tensorflow as tf
+from custom_utils.custom_utils import UV2Mesh, load_model, maskFacesBatch
-from custom_utils.custom_utils import UV2Mesh, maskFacesBatch
 import plotly.graph_objects as go
 import plotly.express as ex
 
@@ -20,15 +32,16 @@ from custom_utils.smpl import SMPL
 # from models.upscaler import Upscaler
 from models.upscaler import UpscalerBasic as Upscaler
 
-name='up_2_3_d128'
+name='up_2_3_skirt_d32'
 def init_wandb():
     run = wandb.init(project='vae', name=name,
-            notes='convnext upscaling\
+            notes='\
 changed ms-ssim scale weights\
 inpainting style\
 MLP based\
 smaller batch\
 use mask for loss\
+started from up_2_3\
     ')
     artifact = wandb.Artifact('code', type='code')
     artifact.add_file('upscale_test.py')
@@ -38,7 +51,7 @@ use mask for loss\
     artifact.add_file('convnext/models/convnext_tf.py')
     run.log_artifact(artifact)
 
-uv2mesh = UV2Mesh(res=256)
+uv2mesh = UV2Mesh(res=256, is_skirt=True)
 
 
 
@@ -87,9 +100,9 @@ def train_step(batch):
         mesh_pred = uv2mesh(pred)
         mesh_orig = batch['mesh_unposed']
 
-        l1_loss = tf.reduce_mean(tf.reduce_sum(((tf.math.abs(tf.transpose(batch['uv_unposed'], [0,2,3,1]) - pred * batch['uv_mask'][:,:,:,tf.newaxis]) )), axis=1))
+        l1_loss = tf.reduce_mean(tf.reduce_sum(((tf.math.abs(batch['uv_unposed'] - pred * batch['uv_mask'][:,:,:,tf.newaxis]) )), axis=1))
-        ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(pred * batch['uv_mask'][:,:,:,tf.newaxis] + 1, tf.transpose(batch['uv_unposed'], [0,2,3,1]) + 1, max_val=2, power_factors=(0.5, 0.2, 0.1, 0.1, 0.1)))
+        # ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(pred * batch['uv_mask'][:,:,:,tf.newaxis] + 1, tf.transpose(batch['uv_unposed'], [0,2,3,1]) + 1, max_val=2, power_factors=(0.5, 0.2, 0.1, 0.1, 0.1)))
-        # ms_ssim = tf.reduce_mean(tf.image.ssim(tf.transpose(pred * uv_mask_body[tf.newaxis,tf.newaxis,:,:], [0,2,3,1]) + 1, tf.transpose(batch['uv_body_posed' + resolution_modifier] * uv_mask_body[tf.newaxis,tf.newaxis,:,:], [0,2,3,1]) + 1, max_val=2))
+        ms_ssim = tf.reduce_mean(tf.image.ssim(pred * batch['uv_mask'][:,:,:,tf.newaxis] + 1, batch['uv_unposed'] * batch['uv_mask'][:,:,:,tf.newaxis] + 1, max_val=2))
         loss = 1e3 * (1 - ms_ssim)
 
     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'])
@@ -124,7 +137,7 @@ def plot_everything(batch: dict, type: str):
     mesh_pred = uv2mesh(pred)
     mesh_orig = batch['mesh_unposed']
     
-    F = np.load("custom_utils/smpl_faces.npy")
+    F = np.load("custom_utils/smpl_faces_skirt.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])]
@@ -133,7 +146,7 @@ def plot_everything(batch: dict, type: str):
     figs = { type + "/garment_" + str(i): plot_3d(gt[i], garment_unposed[i], faces[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) }
     
-    orig_img = tf.transpose(batch['uv_unposed'], [0,2,3,1])
+    orig_img = batch['uv_unposed']
     img = tf.concat((orig_img, pred), axis=2)
     imgs = [wandb.Image((img[i] + 1) / 2.0) for i in range(4)]
 
@@ -142,7 +155,7 @@ def plot_everything(batch: dict, type: str):
 
 
 def train():
-    valid_iterator = iter(dataset_valid)
+    # valid_iterator = iter(dataset_valid)
 
     step = 0
     for batch in dataset_train:
@@ -180,7 +193,7 @@ if __name__ == '__main__':
     # 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)
 
-    low_res = 128
+    low_res = 32
     uv_mask_body = tf.ones((low_res, low_res//2))
     resolution_modifier = '_' + str(low_res)
 
@@ -197,10 +210,10 @@ if __name__ == '__main__':
                          'uv_mask' + resolution_modifier: (low_res, low_res//2,),
                         }
 
-    dataset_train = factory.get_dataset(folder='/home/workspace/local-data/train_100', batch_size=8, signature=signature_dynamic, use_filled=False, resolutions=[512, low_res])
+    dataset_train = factory.get_dataset(folder='/home/workspace/local-data/train_skirt_25_2', batch_size=8, signature=signature_dynamic, use_filled=False, resolutions=[512, low_res])
-    static_dataset_train = StaticDataset(folder='/home/workspace/local-data/train_100', signature=signature_static, use_filled=False, resolutions=[512, low_res])
+    static_dataset_train = StaticDataset(folder='/home/workspace/local-data/train_skirt_25_2', signature=signature_static, use_filled=False, resolutions=[512, low_res])
-    dataset_valid = factory.get_dataset(folder='/home/workspace/local-data/jumpsuit_20/valid', batch_size=32, signature=signature_dynamic)
+    dataset_valid = None #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)
+    static_dataset_valid = None #StaticDataset(folder='/home/workspace/local-data/jumpsuit_20/valid', signature=signature_static)
 
     model = Upscaler(bottleneck_size=512, size=low_res)
 
@@ -209,6 +222,6 @@ if __name__ == '__main__':
 
     ckpt = tf.train.Checkpoint(model=model)
     manager = tf.train.CheckpointManager(ckpt, 'checkpoints/' + name, max_to_keep=3)
-    ckpt.restore(manager.latest_checkpoint)
+    load_model(model, 'checkpoints/up_2_3')
 
     train()

verify_skirt_uv_maps.py

@@ -0,0 +1,83 @@
+# check if upscaler is better then filling
+import os
+
+from models.upscaler import UpscalerBasic as Upscaler
+os.environ["CUDA_VISIBLE_DEVICES"] = "3"
+
+import tqdm
+import numpy as np
+import plotly.graph_objects as go
+
+from custom_utils.custom_utils import UV2Mesh, _maskFaces, load_model
+# from custom_utils.custom_utils import UV2Mesh, maskFacesBatch, _maskFaces
+
+
+if __name__ == '__main__':
+    # folder = '/home/workspace/data/dataset/skirt_test'
+    folder = '/home/workspace/data/dataset/train_skirt_25_2'
+    mask = np.load('custom_utils/mask_256x512_skirt.npy')
+    F = np.load("custom_utils/smpl_faces_skirt.npy")
+    uv2mesh = UV2Mesh(res=256, is_skirt=True)
+    res = 32
+    res_modifier = '_' + str(res)
+
+
+    ids = os.listdir(os.path.join(folder))
+    visited = set()
+    meshes = []
+    count = 0
+
+    for id in tqdm.tqdm(ids):
+        garment_dir = os.path.join(folder, id)
+
+        if not os.path.isdir(garment_dir):
+            continue
+
+        # if os.path.exists(os.path.join(folder, id, 'gt_uv')):
+        #     continue
+
+        garments = [f.replace('.pc16','') for f in os.listdir(garment_dir) if f.find('_') == -1 and f != 'info.mat' and not f.endswith('.obj')]
+        # assert len (garments) == 1
+        if len (garments) != 1:
+            continue
+        garment = garments[0]
+
+        # id, frame_nr = frame.split('_')
+        # frame_nr = frame_nr.split('.')[0]
+        frame_nr = '0'
+        print(id)
+        uv_unposed_no_shape = np.load(os.path.join(folder, id,  'uv_unposed_no_shape', garment + '_' + frame_nr + '.npz'));  uv_unposed_no_shape = uv_unposed_no_shape['uv_unposed_no_shape']
+        # uv_posed = np.load(os.path.join(folder, id,  'gt_uv', garment + '_' + frame_nr + '.npz'));  uv_posed = uv_posed['uv_gt']
+        mesh_mask = np.load(os.path.join(folder, id, garment + '_mask.npy'))
+        uv_mask = np.load(os.path.join(folder, id, 'mask_uv.npy'))
+
+        # body_posed_uv = np.load(os.path.join(folder, id, 'body_posed_uv', garment + '_' + frame_nr + '.npz'))
+        # body_posed_uv = body_posed_uv['uv_body']
+
+        mesh = uv2mesh(uv_unposed_no_shape[np.newaxis,:,:,:]).numpy()[0]
+        # mesh = uv2mesh(body_posed_uv[np.newaxis,:,:,:]).numpy()[0]
+
+        mesh, faces = _maskFaces(mesh, F, mesh_mask)
+        # faces = F
+        x = mesh
+        meshes = [go.Mesh3d(
+                    x=x[:,0],
+                    y=x[:,2],
+                    z=x[:,1],
+                    i = faces[:,0],
+                    j = faces[:,1],
+                    k = faces[:,2]
+                    ),]
+
+        fig = go.Figure(data=meshes)
+        fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=1, z=1))
+        fig.update_layout(
+            scene = dict(
+                xaxis = dict(nticks=4, range=[-1,1],),
+                            yaxis = dict(nticks=4, range=[-1,1],),
+                            zaxis = dict(nticks=4, range=[-1,1],),)
+                            )
+
+        fig.show()
+
+        print()