remove long edges in post

custom_utils/custom_utils.py

@@ -287,4 +287,37 @@ def write_obj(file, vertices, faces, colors):
             file.write(line)
         for f in faces:
             line = 'f ' + ' '.join([str(_) for _ in f]) + '\n'
-            file.write(line)
+            file.write(line)
+
+
+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_normals' + res_modifier], batch['uv_mask' + res_modifier], batch['uv_mask']])
+    conds = [batch['uv_static'], batch['uv_body_posed'], batch['uv_normals']]
+    return levels_input, conds
+
+
+def remove_long_edges(vertices, faces, limit=0.1):
+    # 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 > limit]
+
+    # 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

data/dataset.py

@@ -80,6 +80,12 @@ class StaticCache():
         mesh_mask = np.load(os.path.join(self.folder, 'mesh_mask', f'{id}_{garment}.npy')); mesh_mask = mesh_mask.astype(bool)
         data['mesh_mask'] = mesh_mask
 
+        mesh_fix = np.load(os.path.join(self.folder, 'mesh_fixed', f'{id}_{garment}.npz'))
+        data['mesh_fix'] = dict()
+        data['mesh_fix']['mesh_mask'] = torch.from_numpy(mesh_fix['mask'])
+        # data['mesh_fix']['mask_convert'] = torch.from_numpy(mesh_fix['mas_convert'])
+        # data['mesh_fix']['faces'] = torch.from_numpy(mesh_fix['faces'])
+
         for res in self.resolutions:
             res_string = ''
             if res != 512:
@@ -107,13 +113,15 @@ class StaticCache():
     def __getitem__(self, frame: Frame) -> dict:
         if (frame.id, frame.garment) not in self.store:
             item = self.read(frame.id, frame.garment)
-            stored_data = { key: torch.from_numpy(value).to('cpu') for key, value in item.items()}
+            stored_data = { key: torch.from_numpy(value).to('cpu') for key, value in item.items()  if type(value) == np.ndarray}
+            stored_data['mesh_fix'] = item['mesh_fix']
             for key, value in stored_data.items():
                 if key.startswith('uv_') and value.shape[-1] == 3:
                     stored_data[key] = torch.permute(value, [2,0,1])
             self.store[(frame.id, frame.garment)] = stored_data
         
-        data = { key: value.to(self.device) for key, value in self.store[(frame.id, frame.garment)].items()}
+        data = { key: value.to(self.device) for key, value in self.store[(frame.id, frame.garment)].items() if type(value) == torch.Tensor }
+        data['mesh_fix'] = { key: value.to(self.device) for key, value in self.store[(frame.id, frame.garment)]['mesh_fix'].items() }
         return data
 
 

data/preprocessing.py

@@ -33,7 +33,7 @@ sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
 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
+from custom_utils.custom_utils import _maskFaces, remove_long_edges
 
 
 def create_extras_file(folder):
@@ -631,7 +631,6 @@ def create_frames_file_partial(folder, out_folder, split=0.9, part=0.1):
                 f.write(f'{frame.id} {frame.frame_nr} {frame.garment}\n')
 
 
-
 def check_irregular_poses(folder, folder_reg, out_folder):
     for i in range(100, 200, 5):
         frame: Frame = Frame('01512', i, 'Dress')
@@ -800,6 +799,83 @@ def get_ids():
     pass
 
 
+def generate_fixed_masks(folder, out_folder):
+    ids = sorted(os.listdir(folder))
+    os.makedirs(os.path.join(out_folder, 'mesh_fixed'), exist_ok=True)
+    F = np.load('smpl/smpl_faces.npy')
+    F_skirt = np.load('smpl/smpl_faces_skirt.npy')
+
+    import random
+    # import trimesh
+    import pyvista as pv
+    random.shuffle(ids)
+
+
+    for id in tqdm(ids):
+        garment_dir = os.path.join(folder, id)
+        if not os.path.isdir(garment_dir):
+            continue
+
+        garments = [f.replace('.obj','') for f in os.listdir(garment_dir) if f.endswith('.obj')]
+
+        for garment in garments:
+            if os.path.exists(os.path.join(out_folder, f'{id}_{garment}.pc16')):
+                continue
+
+            is_skirt = garment in ['Dress', 'Skirt']
+
+            mesh_static = np.load(os.path.join(out_folder, 'mesh_static', f'{id}_{garment}.npy'))
+            mask = np.load(os.path.join(out_folder, 'mesh_mask', f'{id}_{garment}.npy'))
+            _, faces = _maskFaces(np.zeros((14475, 3)), F_skirt if is_skirt else F, mask)
+            _, new_mask = remove_long_edges(mesh_static, faces.astype(np.int32), 0.2 if is_skirt else 0.1)
+            new_mask_full = np.zeros((14475), dtype=bool)
+            new_mask_full[mask] = new_mask
+            _, new_faces = _maskFaces(np.zeros((14475, 3)), F_skirt if is_skirt else F, new_mask_full)
+
+            mesh_new = mesh_static[new_mask]
+
+            # fig = go.Figure(data=[go.Mesh3d(x=mesh_new[:,0], y=mesh_new[:,1], z=mesh_new[:,2], i=new_faces[:,0], j=new_faces[:,1], k=new_faces[:,2])])
+            # fig.update_layout(scene_aspectmode='cube')
+            # fig.update_layout(scene=dict(xaxis=dict(range=[-1,1]), yaxis=dict(range=[-1,1]), zaxis=dict(range=[-1,1])))
+            # fig.show()
+
+            flattened_faces = np.asarray([np.concatenate((np.array([3]), f)) for f in new_faces.astype(np.int32)]).flatten()
+            m = pv.PolyData(mesh_new, flattened_faces)
+
+            m.fill_holes(0.03, inplace=True, progress_bar=False)
+
+            # m = trimesh.Trimesh(vertices=mesh_new, faces=new_faces, process=False, validate=False)
+            # trimesh.repair.fill_holes(m)
+
+            mesh_new = np.asarray(m.points)
+            new_faces = np.asarray(m.faces).reshape([-1, 4])[:, 1:]
+
+            m = tm.Trimesh(vertices=mesh_new, faces=new_faces, process=False, validate=False)
+            tm.repair.fix_winding(m)
+
+            mesh_new = np.asarray(m.vertices)
+            new_faces = np.asarray(m.faces)
+
+            # fig = go.Figure(data=[go.Mesh3d(x=mesh_new[:,0], y=mesh_new[:,1], z=mesh_new[:,2], i=new_faces[:,0], j=new_faces[:,1], k=new_faces[:,2])])
+            # fig.update_layout(scene_aspectmode='cube')
+            # fig.update_layout(scene=dict(xaxis=dict(range=[-1,1]), yaxis=dict(range=[-1,1]), zaxis=dict(range=[-1,1])))
+            # fig.show()
+
+            # fig = go.Figure(data=[go.Mesh3d(x=mesh_static[:,0], y=mesh_static[:,1], z=mesh_static[:,2], i=faces[:,0], j=faces[:,1], k=faces[:,2])])
+            # fig.update_layout(scene_aspectmode='cube')
+            # fig.update_layout(scene=dict(xaxis=dict(range=[-1,1]), yaxis=dict(range=[-1,1]), zaxis=dict(range=[-1,1])))
+            # fig.show()
+
+            # print((~new_mask).sum())
+
+            np.savez_compressed(os.path.join(out_folder, 'mesh_fixed', f'{id}_{garment}.npz'), mask=new_mask_full, mas_convert=new_mask, faces=new_faces)
+
+            pass
+
+
+
+
+
 if __name__ == '__main__':
 
 
@@ -832,9 +908,11 @@ if __name__ == '__main__':
     # calculate_average(folder, out_folder)
 
     # create_frames_file(folder, out_folder)
-    get_ids()
+    # get_ids()
     # create_frames_file_test(folder, out_folder)
     # create_frames_file_partial(folder, out_folder, part=0.5)
     
     # folder = '/data-local/127_2/cape_processed/test/female'
-    # create_mesh_normals_cape(folder, None)
+    # create_mesh_normals_cape(folder, None)
+
+    generate_fixed_masks(folder, out_folder)

scripts/cvae_pyramid.py

@@ -238,8 +238,6 @@ def train_step(batch):
     res_modifier = '_' + str(res) if res != 512 else ''
     pred, mean, logvar = out_levels[i]
 
-
-
     l2_loss = torch.mean(1e3 * torch.mean(torch.sum(((uv_orig - pred) ** 2), dim=-1), dim=[1, 2]))
     kl = kl_loss(mean, logvar)
 
@@ -251,16 +249,16 @@ def train_step(batch):
             gender = torch.greater_equal(shape[:, 10] , 0.5).to(dtype=torch.int64)
             shape_offsets = torch.stack([torch.tensordot(shape[i,:10][None, :], smpl.shapedirs[gender[i]], dims=[[1], [2]])[0,smpl.to_keep,:] for i in range(shape.shape[0])])
 
-            mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_mask'][:,:,None]
+            mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_fix']['mesh_mask'][:,:,None]
             mesh_pred_posed = smpl(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped)[0]
         else:
             mesh_pred_posed = mesh_pred
 
-        e_loss = torch.mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']))
+        e_loss = torch.mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_fix']['mesh_mask']))
-        c_loss, c_ratio = collision_loss(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask'])
+        c_loss, c_ratio = collision_loss(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_fix']['mesh_mask'])
         c_loss, c_ratio = torch.mean(c_loss), torch.mean(c_ratio)
-        n_loss = normal_loss(mesh_pred_posed, batch['mesh_posed'], batch['mesh_mask'], is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape'])
+        n_loss = normal_loss(mesh_pred_posed, batch['mesh_posed'], batch['mesh_fix']['mesh_mask'], is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape'])
-        l2_mesh_loss = torch.mean(1e3 * torch.mean(torch.sum(batch['mesh_mask'][:,:,None] * (batch['mesh_posed'] - mesh_pred_posed) ** 2, dim=-1), dim=-1))
+        l2_mesh_loss = torch.mean(1e3 * torch.mean(torch.sum(batch['mesh_fix']['mesh_mask'][:,:,None] * (batch['mesh_posed'] - mesh_pred_posed) ** 2, dim=-1), dim=-1))
         
         reconstr_loss_total = reconstr_loss_total + config['pyramid']['loss_weights'][i] * (l2_loss + l2_mesh_loss)
         if not config['is_cape']:
@@ -290,19 +288,19 @@ def train_step(batch):
                 gender = torch.greater_equal(shape[:, 10] , 0.5).to(dtype=torch.int64)
                 shape_offsets = torch.stack([torch.tensordot(shape[i,:10][None, :], smpl.shapedirs[gender[i]], dims=[[1], [2]])[0,smpl.to_keep,:] for i in range(shape.shape[0])])
 
-                mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_mask'][:,:,None]
+                mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_fix']['mesh_mask'][:,:,None]
                 mesh_pred_posed = smpl(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped)[0]
             else:
                 mesh_pred_posed = mesh_pred
-                
+
-            e_loss = torch.mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']))
+            e_loss = torch.mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_fix']['mesh_mask']))
-            c_loss, c_ratio = collision_loss(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask'])
+            c_loss, c_ratio = collision_loss(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_fix']['mesh_mask'])
             c_loss, c_ratio = torch.mean(c_loss), torch.mean(c_ratio)
-            l2_mesh_loss = torch.mean(1e3 * torch.mean(torch.sum(batch['mesh_mask'][:,:,None] * (batch['mesh_posed'] - mesh_pred_posed) ** 2, dim=-1), dim=-1))
+            l2_mesh_loss = torch.mean(1e3 * torch.mean(torch.sum(batch['mesh_fix']['mesh_mask'][:,:,None] * (batch['mesh_posed'] - mesh_pred_posed) ** 2, dim=-1), dim=-1))
-            n_loss = normal_loss(mesh_pred_posed, batch['mesh_posed'], batch['mesh_mask'], is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape'])
+            n_loss = normal_loss(mesh_pred_posed, batch['mesh_posed'], batch['mesh_fix']['mesh_mask'], is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape'])
         
         ms_ssim = torch.mean(_MSSSIM(pred + 1, batch['uv_static'] + 1))
-        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_fix']['mesh_mask'][:,:,None] * (mesh_orig - mesh_pred.detach()) ** 2, axis=-1))) / torch.sum(batch['mesh_fix']['mesh_mask'])
 
     output = { **output, **{ 'train/edge_loss' + res_modifier: e_loss, 'train/collision_loss' + res_modifier: c_loss,'train/euclid_error' + res_modifier: metric, 
                                 'train/mean' + res_modifier: mean, 'train/logvar' + res_modifier: logvar, 'train/kl_loss' + res_modifier: kl,
@@ -343,18 +341,18 @@ def valid_step(batch):
             gender = torch.greater_equal(shape[:, 10] , 0.5).to(dtype=torch.int64)
             shape_offsets = torch.stack([torch.tensordot(shape[i,:10][None, :], smpl.shapedirs[gender[i]], dims=[[1], [2]])[0,smpl.to_keep,:] for i in range(shape.shape[0])])
 
-            mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_mask'][:,:,None]
+            mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_fix']['mesh_mask'][:,:,None]
             mesh_pred_posed = smpl(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped)[0]
         else:
             mesh_pred_posed = mesh_pred
 
-        e_loss = torch.mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']))
+        e_loss = torch.mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_fix']['mesh_mask']))
-        c_loss, c_ratio = collision_loss(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask'])
+        c_loss, c_ratio = collision_loss(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_fix']['mesh_mask'])
         c_loss, c_ratio = torch.mean(c_loss), torch.mean(c_ratio)
+        l2_mesh_loss = torch.mean(1e3 * torch.mean(torch.sum(batch['mesh_fix']['mesh_mask'][:,:,None] * (batch['mesh_posed'] - mesh_pred_posed) ** 2, dim=-1), dim=-1))
+        n_loss = normal_loss(mesh_pred_posed, batch['mesh_posed'], batch['mesh_fix']['mesh_mask'], is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape'])
         l2_loss = torch.mean(1e3 * torch.mean(torch.sum(((uv_orig - pred) ** 2), dim=-1), dim=[1, 2]))
-        l2_mesh_loss = torch.mean(1e3 * torch.mean(torch.sum(batch['mesh_mask'][:,:,None] * (batch['mesh_posed'] - mesh_pred_posed) ** 2, dim=-1), dim=-1))
+        
-        n_loss = normal_loss(mesh_pred_posed, batch['mesh_posed'], batch['mesh_mask'], is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape'])
-       
         kl = kl_loss(mean, logvar)
 
         if config['use_extra_losses']:
@@ -373,7 +371,7 @@ def valid_step(batch):
         logvar = torch.mean(logvar)
 
         ms_ssim = torch.mean(_MSSSIM(pred + 1, batch['uv_static'] + 1))
-        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_fix']['mesh_mask'][:,:,None] * (mesh_orig - mesh_pred) ** 2, axis=-1))) / torch.sum(batch['mesh_fix']['mesh_mask'])
 
         output = { **output, **{ 'valid/edge_loss' + res_modifier: e_loss, 'valid/collision_loss' + res_modifier: c_loss,'valid/euclid_error' + res_modifier: metric, 
                                     'valid/mean' + res_modifier: mean, 'valid/logvar' + res_modifier: logvar, 'valid/kl_loss' + res_modifier: kl,
@@ -385,63 +383,6 @@ def valid_step(batch):
     return { **{ 'valid/loss': loss, 'valid/reconstr_loss_total': reconstr_loss_total, 'valid/kl_total': kl_total }, **output }
 
 
-def test_step(batch):
-    levels_input, conds = group_inputs_for_levels(batch, config)
-    out_levels = model([levels_input, conds], training=False)
-
-    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][None, :], smpl.shapedirs[gender[i]], axes=[[1], [2]])[0,:,:], smpl.to_keep, dim=0) for i in range(shape.shape[0])])
-    mesh_orig = batch['mesh_unposed']
-
-    output = dict()
-    reconstr_loss_total = 0
-    kl_total = 0
-
-    #for i, res in enumerate(config['pyramid']['resolutions']):
-    i = -1
-    res = config['pyramid']['resolutions'][-1]
-    res_modifier = '_' + str(res) if res != 512 else ''
-    pred, mean, logvar = out_levels[i]
-
-    mesh_pred = uv2mesh(pred)
-    mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_mask'][:,:,None]
-    mesh_pred_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped)
-
-    # is the output already masked?
-
-    e_loss = torch.mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']))
-    c_loss, c_ratio = collision_penalty(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask'])
-    c_loss, c_ratio = torch.mean(c_loss), torch.mean(c_ratio)
-    l2_loss = torch.mean(1e3 * torch.mean(torch.sum(((batch['uv_unposed'] - pred) ** 2), dim=-1), dim=[1, 2]))
-    l2_mesh_loss = torch.mean(1e3 * torch.mean(torch.sum(batch['mesh_mask'][:,:,None] * (mesh_orig - mesh_pred) ** 2, dim=-1), dim=-1))
-    # ms_ssim = torch.mean(tf.image.ssim_multiscale(pred + 5, batch['uv_unposed'] + 5, max_val=10, power_factors=(0.5, 0.2, 0.1, 0.1, 0.1)))
-    # ms_ssim = torch.mean(tf.image.ssim(pred + 1, batch['uv_unposed'] + 1, max_val=3))
-    kl = kl_loss(mean, logvar)
-
-    reconstr_loss_total = reconstr_loss_total + config['pyramid']['loss_weights'][i] * (l2_loss + l2_mesh_loss)
-    if not config['is_cape']:
-         reconstr_loss_total = reconstr_loss_total + config['pyramid']['loss_weights'][i] * (e_loss + c_loss) * config['reg_loss_coeff']
-    # reconstr_loss_total += config['pyramid']['loss_weights'][i] * (l2_loss + (e_loss + c_loss) * config['reg_loss_coeff'])
-    kl_total += kl
-
-    mean = torch.mean(mean)
-    logvar = torch.mean(logvar)
-
-    metric = torch.sum(tf.math.sqrt(torch.sum(batch['mesh_mask'][:,:,None] * (mesh_orig - mesh_pred) ** 2, dim=-1))) / torch.sum(batch['mesh_mask'])
-    euclid_error_per_garment = torch.sum(tf.math.sqrt(torch.sum(batch['mesh_mask'][:,:,None] * (mesh_orig - mesh_pred) ** 2, dim=-1)), dim=1) / torch.sum(batch['mesh_mask'], dim=1)
-
-    output = { **output, **{ 'test/edge_loss' + res_modifier: e_loss, 'test/collision_loss' + res_modifier: c_loss,'test/euclid_error' + res_modifier: metric, 
-                                'test/mean' + res_modifier: mean, 'test/logvar' + res_modifier: logvar, 'test/kl_loss' + res_modifier: kl,
-                                'test/var' + res_modifier: np.e ** logvar, 'test/l2_loss' + res_modifier: l2_loss,
-                                'test/euclid_error': metric, 'test/collision_ratio' + res_modifier: c_ratio } }
-
-    loss = reconstr_loss_total + balancing_coeff * kl_total
-
-    return { **{ 'test/loss': loss, 'test/reconstr_loss_total': reconstr_loss_total, 'test/kl_total': kl_total }, **output }, euclid_error_per_garment
-
-
-
 def plot_everything(batch: dict, type: str):
     if config['use_skinning']:
         mesh_orig, uv_orig = batch['mesh_unposed'].cpu().numpy(), batch['uv_unposed']
@@ -473,8 +414,11 @@ def plot_everything(batch: dict, type: str):
     imgs = [wandb.Image(img[i].cpu().numpy()) for i in range(4)]
 
     _F = np.load("smpl/smpl_faces" + skirt_modifier + ("_orig" if config['is_cape'] else "") + ".npy")
-    mask = batch['mesh_mask'].cpu().numpy().astype(bool)
+    mask = batch['mesh_fix']['mesh_mask'].cpu().numpy().astype(bool)
-    _, faces = maskFacesBatch(np.ones_like(mesh_orig), _F, mask)
+    if 'mesh_fix' in batch:
+        faces = [np.load( (dataset_train if type == 'train' else dataset_valid).folder + "/mesh_fixed/" + batch['id'][i] + '_' + batch['garment'][i] + ".npz" )['faces'] for i in range(len(batch['id']))]
+    else:
+        _, faces = maskFacesBatch(np.ones_like(mesh_orig), _F, mask)
     
     if config['use_skinning']:
         shape = batch['pose_shape'][:, 72:]
@@ -543,237 +487,6 @@ def plot_everything(batch: dict, type: str):
     return { **figs, **figs_coll, **{ type + '/images': imgs, }, **{ f'{type}/batch_info': batch_info }, }
 
 
-def sample(folder):
-    # run = wandb.init(project='vae_plots', name=folder.split('/')[-1])
-    load_model(model,folder)
-    count = 0
-    for batch in dataset_train:
-        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
-
-        pred = model.sample(batch['uv_static'][0], batch['uv_body_posed'][0], uv_curvature[0], batch['uv_mask'][0], nr_of_samples=nr_of_samples)
-        pose, shape = batch['pose_shape'][0, :72], batch['pose_shape'][0, 72:]
-        pose, shape = tf.tile(pose[None,:], [nr_of_samples,1]), tf.tile(shape[None,:], [nr_of_samples,1])
-
-        mesh_pred = uv2mesh(tf.transpose(pred, [0,3,1,2]))
-        mesh_posed = smpl.run_garment(pose, shape, mesh_pred).numpy()
-        mesh_orig = batch['mesh_posed'][0].numpy()
-        mesh_orig_reposed = smpl.run_garment(pose[:1], shape[:1], batch['mesh_unposed'][:1]).numpy()[0]
-
-        body = batch['mesh_body_posed'][0].numpy()
-        body_color = np.ones_like(body, dtype=np.int32)*100
-
-        F = np.load("custom_utils/smpl_faces.npy")
-        mask = batch['mesh_mask'].numpy().astype(bool)[0]
-        _, faces = _maskFaces(np.ones((14475, 3)), F, mask)
-
-
-        pred_zero = model.sample(batch['uv_static'][0], batch['uv_body_posed'][0], uv_curvature[0], batch['uv_mask'][0], nr_of_samples=nr_of_samples, zero_latent=True)
-        mesh_pred_zero = uv2mesh(tf.transpose(pred_zero, [0,3,1,2]))
-        mesh_posed_zero = smpl.run_garment(pose, shape, mesh_pred_zero).numpy()
-
-
-
-
-        meshes = []
-        meshes.append(go.Mesh3d(
-                    # 8 vertices of a cube
-                    x=mesh_orig_reposed[mask,0],
-                    y=mesh_orig_reposed[mask,2],
-                    z=mesh_orig_reposed[mask,1],
-                    # i, j and k give the vertices of triangles
-                    i = faces[:,0],
-                    j = faces[:,1],
-                    k = faces[:,2]
-                    ))
-
-        for i in range(nr_of_samples):
-            x = mesh_posed[i][mask, :].copy()
-            # x[:, 2] += 1 * (i + 1)
-            x[:, 2] += 1
-            meshes.append(go.Mesh3d(
-                        # 8 vertices of a cube
-                        x=x[:,0],
-                        y=x[:,2],
-                        z=x[:,1],
-                        # i, j and k give the vertices of triangles
-                        i = faces[:,0],
-                        j = faces[:,1],
-                        k = faces[:,2]
-                        ))
-
-        
-        for i in range(nr_of_samples):
-            x = mesh_posed_zero[i][mask, :].copy()
-            # x[:, 2] += 1 * (i + 1)
-            x[:, 2] += 2
-            meshes.append(go.Mesh3d(
-                        # 8 vertices of a cube
-                        x=x[:,0],
-                        y=x[:,2],
-                        z=x[:,1],
-                        # i, j and k give the vertices of triangles
-                        i = faces[:,0],
-                        j = faces[:,1],
-                        k = faces[:,2]
-                        ))
-
-        # for i in range(5):
-        #     x = body.copy()
-        #     x[:, 2] += 1 * i
-        #     meshes.append(go.Mesh3d(
-        #                 # 8 vertices of a cube
-        #                 x=x[:,0],
-        #                 y=x[:,2],
-        #                 z=x[:,1],
-        #                 # i, j and k give the vertices of triangles
-        #                 i = F[:,0],
-        #                 j = F[:,1],
-        #                 k = F[:,2],
-        #                 vertexcolor=body_color
-        #                 ))
-
-
-        fig = go.Figure(data=meshes)
-
-        fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=2, z=1))
-        # fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=nr_of_samples-1, z=1))
-        fig.update_layout(
-            scene = dict(
-                xaxis = dict(nticks=4, range=[-1,1],),
-                            # yaxis = dict(nticks=4, range=[-1,nr_of_samples+1],),
-                            yaxis = dict(nticks=4, range=[-1,3],),
-                            zaxis = dict(nticks=4, range=[-1,1],),)
-                            )
-
-        # wandb.log({'plot_' + str(count): fig})
-        fig.show()
-
-        pred = pred.numpy()
-        frames = [Image.fromarray(((pred[i] + 1) * 128).astype(np.int8), mode='RGB') for i in range(pred.shape[0])]
-        frames[0].save("temp/garment_1.gif", format="GIF", append_images=frames[1:],
-               save_all=True, duration=500, loop=0)
-
-
-        if count > 10:
-            break
-
-        print(count)
-        count += 1
-
-        print()
-
-
-def test():
-    garment_type_table = np.load('custom_data/garment_type_table_test.npy', allow_pickle=True).item()
-    str_split = config['name'].split('_')
-    test_idx = str_split.index('test')
-    checkpoint_name = '_'.join(str_split[:test_idx])
-    load_model(model, 'checkpoints/' + checkpoint_name)
-
-    dataset_test = factory.get_dataset(folder=config['test_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'], use_filled=config['use_filled'], resolutions=dataset_resolutions, in_memory=False, do_repeat=False)
-    static_dataset_test = StaticDataset(folder=config['test_dir'], signature=config['signature_static'], use_filled=config['use_filled'], resolutions=dataset_resolutions, in_memory=False)
-   
-    step = 0
-    losses = []
-
-    if config['is_skirt']:
-        per_garment_errors = { 'Tshirt': [], 'Dress': [], 'Skirt': [], 'Top': [] }
-    else:
-        per_garment_errors = { 'Tshirt': [], 'Jumpsuit': [], 'Trousers': [], 'Top': [] }
-
-    for batch in tqdm.tqdm(dataset_test):
-        batch = static_dataset_test.expand_batch(batch_to_dict(batch, config['signature_dynamic']))
-
-        test_loss, euclid_error_per_garment = test_step(batch)
-        losses.append(test_loss)
-
-        for i, id in enumerate(batch['id']):
-            id_str = id.numpy().decode("utf-8")
-            per_garment_errors[garment_type_table[id_str]].append(euclid_error_per_garment[i].numpy()) 
-
-        to_log = { **{ 'current_step': step }, **test_loss }
-        if step % 1000 == 0:
-            figs = plot_everything(batch, 'test')
-            to_log = { **to_log, **figs }
-
-        wandb.log(to_log, step = step)
-        
-        # print(step)
-        step += 1
-   
-    # averages = dict()
-    path_to_plotly_html = "temp/plotly_figure.html"
-
-    histograms = dict()
-    for key in test_loss.keys():
-        accumulated = [losses[i][key].numpy() for i in range(len(losses))]
-        fig = ex.histogram(accumulated, title='Average: ' + str(np.asarray(accumulated).mean()))
-        fig.write_html(path_to_plotly_html, auto_play = False)
-
-        histograms['all/' + key.split('/')[-1] + '_hist'] = wandb.Html(path_to_plotly_html)
-        # averages['all/' + key.split('/')[-1] + '_avg'] = np.asarray(accumulated).mean()
-
-    for key in per_garment_errors.keys():
-        fig = ex.histogram(per_garment_errors[key], title='Average: ' + str(np.asarray(per_garment_errors[key]).mean()))
-        fig.write_html(path_to_plotly_html, auto_play = False)
-        histograms['all/' + key + 'euclid_error_hist'] = wandb.Html(path_to_plotly_html)
-    
-
-    wandb.log({ **histograms }, step = step)
-    wandb.finish()
-    pass
-
-
-def generate_results():
-    load_model(model, 'checkpoints/' + config['name'])
-
-    dataset_test = factory.get_dataset(folder=config['test_dir'], 
-                                        batch_size=config['batch_size'], 
-                                        # batch_size=1, 
-                                        signature=config['signature_dynamic'], 
-                                        resolutions=dataset_resolutions, 
-                                        in_memory=config['dataset_in_memory'], 
-                                        is_cape=config['is_cape'],
-                                        frames_file=config['frames_file'] if 'frames_file' in config.keys() else None,
-                                        do_repeat=False,
-                                        shuffle=False)
-    static_dataset_test = StaticDataset(folder=config['test_dir'], 
-                                        signature=config['signature_static'], 
-                                        resolutions=dataset_resolutions, 
-                                        in_memory=True,
-                                        is_cape=config['is_cape'])
-    folder = os.path.join('/home/workspace/data/results/cape', config['name'])
-    os.makedirs(folder, exist_ok=True)
-   
-    for batch in tqdm.tqdm(dataset_test):
-        batch = static_dataset_test.expand_batch(batch_to_dict(batch, config['signature_dynamic']))
-
-        levels_input, conds = group_inputs_for_levels(batch, config)
-        out_levels = model([levels_input, conds], training=False)
-
-        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][None, :], smpl.shapedirs[gender[i]], axes=[[1], [2]])[0,:,:], smpl.to_keep, axis=0) for i in range(shape.shape[0])])
-
-        i = -1
-        pred = out_levels[i][0]
-
-        mesh_pred = uv2mesh(pred)
-        mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_mask'][:,:,None]
-        mesh_pred_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped)
-
-        for i in range(batch['id'].shape[0]):
-            mask = batch['mesh_mask'][i].numpy().astype(bool)
-            id, frame_nr = batch['id'][i].numpy().decode("utf-8"), batch['frame_nr'][i].numpy().decode("utf-8")
-            unposed = mesh_pred[i].numpy()[mask, :]
-            posed = mesh_pred_posed[i].numpy()[mask, :]
-            np.save(os.path.join(folder, id + '_' + frame_nr + '_posed.npy'), posed)
-            np.save(os.path.join(folder, id + '_' + frame_nr + '_unposed.npy'), unposed)
-
-    pass
-
 def speed_step(levels_input, conds):
 
     return model.sample([levels_input, conds], training=False)
@@ -840,7 +553,7 @@ if __name__ == '__main__':
 
     optimizer = optim.Adam(model.parameters(), lr=config['lr'])
     _MSSSIM = MS_SSIM(data_range=2, size_average=True, channel=3)
-    ckpt_saver = CheckpointSaver('checkpoints/iccvw/' + config['name'], top_n=5)
+    ckpt_saver = CheckpointSaver('checkpoints/iccvw/' + config['name'], top_n=20)
     scaler = torch.cuda.amp.GradScaler(enabled=use_amp)
 
     g_balancing_coeff = config['starting_balancing_coeff']