investigate duplicate vertices

.gitignore

@@ -172,3 +172,4 @@ ids/
 examples/
 custom_data/dijkstra*
 .vscode/*
+.vscode/launch.json

.vscode/launch.json

@@ -10,7 +10,7 @@
             "request": "launch",
             "program": "${file}",
             "console": "integratedTerminal",
-            "justMyCode": true,
+            "justMyCode": false, // true false
             "cwd": "${workspaceFolder}"
         }
     ]

compression_test.py

@@ -0,0 +1,32 @@
+import compress_pickle as cp
+import numpy as np
+import os
+
+def convert_bytes(num):
+    """
+    this function will convert bytes to MB.... GB... etc
+    """
+    for x in ['bytes', 'KB', 'MB', 'GB', 'TB']:
+        if num < 1024.0:
+            return "%3.1f %s" % (num, x)
+        num /= 1024.0
+
+
+def file_size(file_path):
+    """
+    this function will return the file size
+    """
+    if os.path.isfile(file_path):
+        file_info = os.stat(file_path)
+        return convert_bytes(file_info.st_size)
+    
+
+if __name__ == '__main__':
+    # Create a random array
+    uv_mask = np.load('custom_utils/mask_256x512.npy')
+    arr = np.stack([np.random.rand(512, 256, 3) * uv_mask[:,:,None] for _ in range(5)])
+    cp.dump(arr, 'temp/test_pickle', 'bz2')
+    np.savez_compressed('temp/test_numpy.npz', arr)
+
+    print('Pickle file size: ', file_size('temp/test_pickle.' + 'bz'))
+    print('Numpy file size: ', file_size('temp/test_numpy.npz'))

data/cloth3d_io.py

@@ -231,10 +231,6 @@ def read_posed_garment_frame(folder, id, garment, frame, map=None, mask=None, sh
         map = sio.loadmat(os.path.join(folder, id, garment + '_map.mat'))['map'].reshape((-1))
     verts = verts[map, :]
 
-    if unshape:
-        shape_offsets = np.tensordot(shape, shapedirs, axes=[[1], [2]])
-        verts[mask, :] = verts[mask, :] - np.squeeze(shape_offsets)[mask, :]
-
     if deorient:
         R = scipy.spatial.transform.Rotation.from_euler('xyz', pose[:3]).inv()
         # R = scipy.spatial.transform.Rotation.from_euler([0,0,-zrot])
@@ -244,22 +240,17 @@ def read_posed_garment_frame(folder, id, garment, frame, map=None, mask=None, sh
 
 
 # can i cache map?
-# NOT WORKING WELL
 def read_posed_garment(folder, id, garment, map=None, mask=None, shape=None, shapedirs=None, poses=None, unshape=False, deorient=False):
     # body_posed = readPC2(os.path.join(folder, id, 'Body_posed.pc16'), True)['V']
     # offsets = readPC2(os.path.join(folder, id, garment + '_offsets_dynamic.pc16'), True)['V']
     # verts = (body_posed + offsets) * mask[:, None]
     pc16_path = os.path.join(folder, id, garment + '.pc16')
-    info = loadInfo(os.path.join(folder, id, 'info.mat'))
-    verts = np.stack([readPC2Frame(pc16_path, frame, True) for frame in range(info['poses'].shape[1])])
+    # info = loadInfo(os.path.join(folder, id, 'info.mat'))
+    verts = readPC2(pc16_path, True)['V']
     if map is None:
         map = sio.loadmat(os.path.join(folder, id, garment + '_map.mat'))['map'].reshape((-1))
     verts = verts[:, map, :]
 
-    if unshape:
-        shape_offsets = np.tensordot(shape, shapedirs, axes=[[1], [2]])
-        verts[:, mask, :] = verts[:, mask, :] - np.squeeze(shape_offsets)[None, mask, :]
-
     if deorient:
         for i in range(verts.shape[0]):
             R = scipy.spatial.transform.Rotation.from_euler('xyz', poses[:3, i]).inv()
@@ -298,7 +289,7 @@ def read_posed_body_frame(folder, id, frame, pose=None, deorient=False):
 
 def read_posed_body(folder, id, poses=None, deorient=False):
     pc16_path = os.path.join(folder, id, 'Body_posed.pc16')
-    verts = readPC2(pc16_path, True)
+    verts = readPC2(pc16_path, True)['V']
 
     if deorient:
         for i in range(verts.shape[0]):
@@ -309,17 +300,40 @@ def read_posed_body(folder, id, poses=None, deorient=False):
     return verts
 
 
-def read_curvatures(h5file, id, frames, neighbourhood=None):
-    frames_filtered = [frame for frame in frames if frame.split()[0] == id]
-    garment = frames_filtered[0].split()[2]
-    curvatures_all = []
-    for frame_nr in range(len(frames_filtered)):
-        curv = h5file.get_node(f'/{id}/{frame_nr}/{garment}/curvature_raw').read()
-        curvatures_all.append(curv)
-    return np.array(curvatures_all)
+# def read_curvatures(h5file, id, garment, frames, neighbourhood=None):
+#     frames_filtered = [frame for frame in frames if (frame.split()[0] == id and frame.split()[2] == garment)]
+#     if len(frames_filtered) == 0:
+#         print(f'No frames found for {id} {garment}')
+#     if len(frames_filtered[0].split()) < 3:
+#         print(f'No garment specified for {id} {garment}')
+#         print(frames_filtered[0].split())
+#     garment = frames_filtered[0].split()[2]
+#     curvatures_all = []
+#     for frame_nr in range(len(frames_filtered)):
+#         curv = h5file.get_node(f'/{id}/{frame_nr}/{garment}/curvature_raw').read()
+#         curvatures_all.append(curv)
+#     return np.array(curvatures_all)
 
+def read_curvatures(folder, id, garment, frames=None, count=None, neighbourhood=None, per_garment=False):
+    if neighbourhood is not None:
+        raise NotImplementedError('Neighbourhood not implemented for reading curvatures')
+    if count is None:
+        frames_filtered = [frame for frame in frames if (frame.split()[0] == id and frame.split()[2] == garment)]
+        count = len(frames_filtered)
 
+    all = []
+    for i in range(count):
+        curv = np.load(os.path.join(folder, 'curvature_mesh_global' if not per_garment else 'curvature_mesh_garment', f'{id}_{i}_{garment}.npz'))['mesh_curvature_raw'] # not actually raw
+        all.append(curv)
 
+    return all
+
+def read_curvature_frame(folder, id, garment, frame, neighbourhood=None, per_garment=False):
+    if neighbourhood is not None:
+        raise NotImplementedError('Neighbourhood not implemented for reading curvatures')
+    curv = np.load(os.path.join(folder, 'curvature_mesh_global' if not per_garment else 'curvature_mesh_garment', f'{id}_{frame}_{garment}.npz'))['mesh_curvature_raw'] # not actually raw
+
+    return curv
 
 import sys
 sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')

data/convert_to_uv.py

@@ -96,6 +96,73 @@ def convert_to_uv(res, m2uv, uv_map, verts, neighbours=10, mask=None, method='rb
         return result 
 
 
+def convert_to_uv_batch_body(res, m2uv, uv_map, verts, neighbours=15, mask=None, method='rbf', garment_uv_mask=None, garment_uv_mask_border=None):
+    batch = verts.shape[0]
+    verts_uv = np.zeros((batch, res * 2, res, 3))
+    uv_values = np.zeros((batch, uv_map.shape[0], 3))
+    uv_values[:, m2uv[:, 0]] = verts
+    uv_values[:, m2uv[:, 1]] += verts
+    uv_values = uv_values[:, uv_map_mask, :]
+    uv_values /= uv_counter[np.newaxis, :, np.newaxis]
+    return_garment_uv_mask = False
+    # test
+    verts_uv[:, uv_map_discreet_unique[:, 0], uv_map_discreet_unique[:, 1], :] = uv_values
+
+    if mask is not None:
+        garment_mask = np.zeros((uv_map.shape[0]), dtype=bool)
+        garment_mask[m2uv[:, 0]] = mask
+        garment_mask[m2uv[:, 1]] += mask
+        garment_mask = garment_mask[uv_map_mask] #   add garment mask
+    else:
+        garment_mask = np.sum(uv_values[0,:,:], axis=-1) != 0
+
+    unique_uv_values = uv_values[:, garment_mask, :]
+    uv_map_discreet_unique_unique = uv_map_discreet_unique[garment_mask, :]
+    rbf_dest_current = rbf_dest
+
+    # add back mesh vertices that are not part of the garment
+    uv_map_discreet_unique_unique_inv = uv_map_discreet_unique[np.invert(garment_mask), :]
+    rbf_dest_fill = np.zeros((res * 2, res), dtype=bool)
+    rbf_dest_fill[uv_map_discreet_unique_unique_inv[:, 0], uv_map_discreet_unique_unique_inv[:, 1]] = True
+    
+    if garment_uv_mask is None and mask is not None:
+        garment_uv_mask = np.zeros((res * 2, res))
+        garment_uv_mask[uv_map_discreet_unique_unique[:, 0], uv_map_discreet_unique_unique[:, 1]] = 255
+        garment_uv_mask = cv2.morphologyEx(garment_uv_mask, cv2.MORPH_CLOSE, cv2.getStructuringElement(cv2.MORPH_RECT, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=4)
+        garment_uv_mask = cv2.dilate(garment_uv_mask, cv2.getStructuringElement(cv2.MORPH_RECT, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=2).astype(bool)
+        return_garment_uv_mask = True
+
+    if garment_uv_mask is not None:
+        rbf_dest_current = (rbf_dest + rbf_dest_fill) * garment_uv_mask + garment_uv_mask_border
+
+    rbf_dest_current = np.nonzero(rbf_dest_current)
+    rbf_dest_current = np.array(rbf_dest_current).T
+
+    final = []
+    for i in range(batch):
+        print('processing frame: ' + str(i))
+        if method == 'rbf':
+            interpolator =  spi.RBFInterpolator(y=uv_map_discreet_unique_unique, d=unique_uv_values[i,:,:], neighbors=neighbours, kernel='linear')
+        elif method == 'nearest':
+            interpolator =  spi.NearestNDInterpolator(x=uv_map_discreet_unique_unique, y=unique_uv_values[i,:,:])
+        else:
+            raise NotImplementedError
+
+        result = np.zeros((res * 2, res, 3), np.float32)
+        result[rbf_dest_current[:, 0], rbf_dest_current[:, 1], :] = interpolator(rbf_dest_current)
+        result[uv_map_discreet_unique_unique[:, 0], uv_map_discreet_unique_unique[:, 1], :] = unique_uv_values[i,:,:]
+
+        if garment_uv_mask is not None:
+            result *= garment_uv_mask[:,:,np.newaxis]
+        
+        final.append(result)
+
+    if return_garment_uv_mask:
+        return final, garment_uv_mask
+    else:
+        return final 
+    
+
 def convert_to_uv_batch(res, m2uv, uv_map, verts, neighbours=15, mask=None, method='rbf', garment_uv_mask=None, garment_uv_mask_border=None, id=None):
     batch = len(verts)
 
@@ -333,8 +400,8 @@ def get_shapedirs():
 if __name__ == '__main__':
     parser = argparse.ArgumentParser()
     parser.add_argument("-f", "--folder", help = "end of range", default='/home/workspace/data/dataset/cloth3d/train')
-    parser.add_argument("-s", "--start", help = "start of range", default=0)
-    parser.add_argument("-e", "--end", help = "end of range", default=1000)
+    parser.add_argument("-s", "--start", help = "start of range", default=70500)
+    parser.add_argument("-e", "--end", help = "end of range", default=90000)
     parser.add_argument("-n", "--name", help = "end of range", default='lhunor')
     parser.add_argument("-sk", "--skirt", help = "end of range", default=0)
     args = parser.parse_args()
@@ -388,15 +455,25 @@ if __name__ == '__main__':
     with open(os.path.join(folder, frames_file), 'r') as f:
         lines = f.readlines()
         lines = [line.strip() for line in lines]
-    frames = lines
+    frames = lines[int(args.start):int(args.end)]
     ids = [frame.split()[0] for frame in frames]
-    ids = sorted(list(set(ids)))[int(args.start):int(args.end)]
+    ids = sorted(list(set(ids)))
 
     shapedirs = get_shapedirs()
 
     root = '/'.join(folder.split('/')[:-1])
     type = folder.split('/')[-1]
-    h5file = tb.open_file(os.path.join(root, f'{type}_extras.h5'), mode='a')
+    # h5file_curv = tb.open_file(os.path.join(root, f'{type}_extras.h5'), mode='r')
+    # h5file = tb.open_file(os.path.join(root, f'{type}_extras_uv.h5'), mode='a')
+
+    extras_folder = os.path.join(root, f'{type}_extras')
+    os.makedirs(extras_folder, exist_ok=True)
+    os.makedirs(os.path.join(extras_folder, 'uv_posed' + orient_mod), exist_ok=True)
+    os.makedirs(os.path.join(extras_folder, 'uv_static'), exist_ok=True)
+    os.makedirs(os.path.join(extras_folder, 'uv_curvature'), exist_ok=True)
+    os.makedirs(os.path.join(extras_folder, 'uv_unposed'), exist_ok=True)
+    os.makedirs(os.path.join(extras_folder, 'uv_body_posed' + orient_mod), exist_ok=True)
+    os.makedirs(os.path.join(extras_folder, 'uv_mask'), exist_ok=True)
 
 
     for id in tqdm.tqdm(ids):
@@ -411,121 +488,73 @@ if __name__ == '__main__':
         # also masking is not the same as just filling the empty values with neighbouring
         # values, at least the plots are different
 
-        garment = garments[0]
+        for garment in garments:
+        # garment = garments[0]
 
-        mask = read_garment_mask(folder, id, garment)
+            # if all the uv maps have been generated, skip this id
+            # if all([os.path.exists(os.path.join(extras_folder, 'uv_body_posed' + orient_mod,  f'{id}_{i}_{garment}.npz')) for i in range(info['poses'].shape[1])]):
+            #     continue
 
-        if 'mesh_posed' in data_to_convert:
-            mesh_posed_frame = read_posed_garment_frame(folder, id, garment, 0, mask=mask, shape=info['shape'], pose=info['poses'][:,0], deorient=use_no_orient, unshape=no_shape)
-            mesh_posed = read_posed_garment(folder, id, garment, mask=mask, shape=info['shape'], poses=info['poses'], deorient=use_no_orient, unshape=no_shape)
-        if 'mesh_unposed' in data_to_convert:
-            mesh_unposed = read_unposed_garment(folder, id, garment, mask=mask, shape=info['shape'], unshape=no_shape)
-        if 'mesh_static' in data_to_convert:
-            mesh_static = read_garment_template(folder, id, garment, mask)
-        if 'mesh_curvature' in data_to_convert:
-            mesh_curvature = read_curvatures(h5file, id, frames, neighbourhood=None)
-        if 'body_posed' in data_to_convert:
-            body_posed = read_posed_body(folder, id, info['poses'], use_no_orient)
+            mask = read_garment_mask(folder, id, garment)
+
+            if 'mesh_posed' in data_to_convert:
+                # mesh_posed_frame = read_posed_garment_frame(folder, id, garment, 0, mask=mask, shape=info['shape'], pose=info['poses'][:,0], deorient=use_no_orient, unshape=no_shape)
+                mesh_posed = read_posed_garment(folder, id, garment, mask=mask, shape=info['shape'], poses=info['poses'], deorient=use_no_orient, unshape=no_shape)
+            if 'mesh_unposed' in data_to_convert:
+                raise NotImplementedError()
+                mesh_unposed = read_unposed_garment(folder, id, garment, mask=mask, shape=info['shape'], unshape=no_shape)
+            if 'mesh_static' in data_to_convert:
+                mesh_static = read_garment_template(folder, id, garment, mask)
+            if 'mesh_curvature' in data_to_convert:
+                mesh_curvature = read_curvatures(folder + '_extras', id, garment, count=info['poses'].shape[1])
+            if 'body_posed' in data_to_convert:
+                body_posed = read_posed_body(folder, id, info['poses'], use_no_orient)
 
 
-        # verts_static = read_garment_template(folder, id, garment, mask)
-        # body = read_posed_body(folder, id, info['poses'], use_no_orient)
-        # gt = read_posed_garment(folder, id, garment, mask=mask, shape=info['shape'], ...)
-        # gt = readPC2(os.path.join(folder, id, garment + '_registered.pc16'),  True)['V']
-        # curvature = np.stack([np.load(os.path.join(folder, f'mesh_curvature_{neighbourhood}_2',  f'{id}_{i}.npz'))['mesh_curvature'] for i in range(frames_len)])
-        # 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_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])]
-        # body_no_orient = np.stack([body_no_orient[i]['mesh_posed_body_no_orient'] for i in range(gt.shape[0])])
-        # mesh_unposed_no_shape_one_hot = [np.load(os.path.join(local_folder, 'mesh_unposed_no_shape_one_hot',  id + '_' + str(i) + '.npz')) for i in range(frames)]
-        # mesh_unposed_no_shape_one_hot = [mesh_unposed_no_shape_one_hot[i]['mesh_unposed_no_shape'] for i in range(frames)]
-        # mesh_posed_no_orient = [np.load(os.path.join(local_folder, 'mesh_posed_no_orient',  id + '_' + str(i) + '.npz')) for i in range(frames)]
-        # mesh_posed_no_orient = [mesh_posed_no_orient[i]['mesh_posed_no_orient'] for i in range(frames)]
+            start = time.perf_counter()
+
+            uv_static, uv_mask_local, garment_uv_mask_border = convert_to_uv(uv_res, m2uv, uv_map, mesh_static, neighbours=3, mask=mask, method='rbf', garment_uv_mask=None)
+            uv_static = convert_to_uv_batch(uv_res, m2uv, uv_map, mesh_static[None, :, :], neighbours=3, mask=mask, method='rbf', garment_uv_mask=uv_mask_local, garment_uv_mask_border=garment_uv_mask_border)
+            uv_static = uv_static.astype(np.float32)
+
+            if 'mesh_posed' in data_to_convert:
+                uv_posed = convert_to_uv_batch(uv_res, m2uv, uv_map, mesh_posed[:, :, :], neighbours=3, mask=mask, method='rbf', garment_uv_mask=uv_mask_local, garment_uv_mask_border=garment_uv_mask_border)
+            if 'mesh_unposed' in data_to_convert:
+                raise NotImplementedError()
+                uv_unposed = convert_to_uv_batch(uv_res, m2uv, uv_map, mesh_unposed[:, :, :], neighbours=3, mask=mask, method='rbf', garment_uv_mask=uv_mask_local, garment_uv_mask_border=garment_uv_mask_border)
+            if 'mesh_curvature' in data_to_convert:
+                uv_curvature = convert_to_uv_batch_label(uv_res, m2uv, uv_map, mesh_curvature[:, :], neighbours=3, mask=mask, method='nearest', garment_uv_mask=uv_mask_local, garment_uv_mask_border=garment_uv_mask_border)
+            if 'body_posed' in data_to_convert:
+                uv_body_posed = convert_to_uv_batch_body(uv_res, m2uv, uv_map, body_posed[:, :, :], neighbours=3, mask=None, method='rbf', garment_uv_mask=None, garment_uv_mask_border=garment_uv_mask_border)
 
 
-        # verts_dynamic = gt
 
-        # TODO: check if have to mask out non garment vertices or not? are they the same vertices erpeated many times? if remove, have to interpolate from almost nothing?
+            print(time.perf_counter() - start)
 
-        # test(data.f.Vt, data.f.Ft, verts)
-        # test(uv_map, data.f.Ft, verts)
+            # mesh_static
+            # save_to_h5file(h5file, id, 0, garment, uv_static, 'uv_static', 'float32')
+            # save_to_h5file(h5file, id, 0, garment, uv_mask, 'uv_mask', 'float32')
+            np.savez_compressed(os.path.join(extras_folder, 'uv_static',  f'{id}_{garment}.npz'), uv_static=uv_static)
+            np.savez_compressed(os.path.join(extras_folder, 'uv_mask',  f'{id}_{garment}.npz'), uv_mask=uv_mask_local)
 
-        start = time.perf_counter()
+            if 'mesh_posed' in data_to_convert:
+                for i in range(len(uv_posed)):
+                    np.savez_compressed(os.path.join(extras_folder, 'uv_posed' + orient_mod,  f'{id}_{i}_{garment}.npz'), uv_posed=uv_posed[i])
+                    # save_to_h5file(h5file, id, i, garment, uv_posed[i, :, :, :], 'uv_posed' + orient_mod + shape_mod, 'float32')
+            if 'mesh_unposed' in data_to_convert:
+                raise NotImplementedError()
+                for i in range(uv_unposed.shape[0]):
+                    save_to_h5file(h5file, id, i, garment, uv_unposed[i, :, :, :], 'uv_unposed' + orient_mod + shape_mod, 'float32')
+            if 'mesh_curvature' in data_to_convert:
+                for i in range(len(uv_curvature)):
+                    np.savez_compressed(os.path.join(extras_folder, 'uv_curvature',  f'{id}_{i}_{garment}.npz'), uv_curvature=uv_curvature[i].astype(np.uint8))
+                    # save_to_h5file(h5file, id, i, garment, uv_curvature[i, :, :, :], 'uv_curvature', 'uint8')
+            if 'body_posed' in data_to_convert:
+                for i in range(len(uv_body_posed)):
+                    np.savez_compressed(os.path.join(extras_folder, 'uv_body_posed' + orient_mod,  f'{id}_{i}_{garment}.npz'), uv_body_posed=uv_body_posed[i])
+                    # save_to_h5file(h5file, id, i, garment, uv_body_posed[i, :, :, :], 'uv_body_posed', 'float32')
             
-        uv_static, uv_mask, garment_uv_mask_border = convert_to_uv(uv_res, m2uv, uv_map, mesh_static, neighbours=3, mask=mask, method='rbf', garment_uv_mask=None)
-        uv_static = uv_static.astype(np.float32)
-
-        if 'mesh_posed' in data_to_convert:
-            uv_posed = convert_to_uv_batch(uv_res, m2uv, uv_map, mesh_posed[:, :, :], neighbours=3, mask=mask, method='rbf', garment_uv_mask=uv_mask, garment_uv_mask_border=garment_uv_mask_border).astype(np.float32)
-        if 'mesh_unposed' in data_to_convert:
-            uv_unposed = convert_to_uv_batch(uv_res, m2uv, uv_map, mesh_unposed[:, :, :], neighbours=3, mask=mask, method='rbf', garment_uv_mask=uv_mask, garment_uv_mask_border=garment_uv_mask_border).astype(np.float32)
-        if 'mesh_curvature' in data_to_convert:
-            uv_curvature = convert_to_uv_batch_label(uv_res, m2uv, uv_map, mesh_curvature[:, :], neighbours=3, mask=mask, method='nearest', garment_uv_mask=uv_mask, garment_uv_mask_border=garment_uv_mask_border).astype(np.float32)
-        if 'body_posed' in data_to_convert:
-            uv_body_posed = convert_to_uv_batch(uv_res, m2uv, uv_map, body_posed[:, :, :], neighbours=3, mask=None, method='rbf', garment_uv_mask=None, garment_uv_mask_border=garment_uv_mask_border).astype(np.float32)
-
-        print(time.perf_counter() - start)
-
-        # mesh_static
-        save_to_h5file(h5file, id, 0, garment, uv_static, 'uv_static', 'float32')
-        save_to_h5file(h5file, id, 0, garment, uv_mask, 'uv_mask', 'float32')
-        if 'mesh_posed' in data_to_convert:
-            for i in range(uv_posed.shape[0]):
-                save_to_h5file(h5file, id, i, garment, uv_posed[i, :, :, :], 'uv_posed' + orient_mod + shape_mod, 'float32')
-        if 'mesh_unposed' in data_to_convert:
-            for i in range(uv_unposed.shape[0]):
-                save_to_h5file(h5file, id, i, garment, uv_unposed[i, :, :, :], 'uv_unposed' + orient_mod + shape_mod, 'float32')
-        if 'mesh_curvature' in data_to_convert:
-            for i in range(uv_curvature.shape[0]):
-                save_to_h5file(h5file, id, i, garment, uv_curvature[i, :, :, :], 'uv_curvature', 'uint8')
-        if 'body_posed' in data_to_convert:
-            for i in range(uv_body_posed.shape[0]):
-                save_to_h5file(h5file, id, i, garment, uv_body_posed[i, :, :, :], 'uv_body_posed', 'float32')
+            pass
 
 
-        # start = time.perf_counter()
-        # np.save(os.path.join(folder, id, 'body_static_uv.npy'), uv_body_static)
-
-        # os.makedirs(os.path.join(folder, id, 'body_posed_uv'), exist_ok=True)
-        # for i in range(uv_body.shape[0]):
-        #     np.savez_compressed(os.path.join(folder, id, 'body_posed_uv', garment + '_' + str(i) + '.npz'), uv_body=uv_body[i, :, :, :])
-        
-        # os.makedirs(os.path.join(folder, id, 'uv_posed_body_no_orient'), exist_ok=True)
-        # 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)
-        # 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, :, :, :])
-
-        os.makedirs(os.path.join(folder, f'uv_curvature_{neighbourhood}_2'), exist_ok=True)
-        for i in range(uv_curvature.shape[0]):
-            np.savez_compressed(os.path.join(folder, f'uv_curvature_{neighbourhood}_2', id + '_' + str(i) + '.npz'), curvature_uv=uv_curvature[i, :, :])
-
-        # os.makedirs(os.path.join(folder, id, 'uv_unposed'), exist_ok=True)
-        # 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, :, :, :])
-
-        # os.makedirs(os.path.join(folder, id, 'uv_unposed_no_shape'), exist_ok=True)
-        # for i in range(uv_unposed_no_shape.shape[0]):
-        #     np.savez_compressed(os.path.join(folder, id, 'uv_unposed_no_shape', garment + '_' + str(i) + '.npz'), uv_unposed_no_shape=uv_unposed_no_shape[i, :, :, :])
-
-        # os.makedirs(os.path.join(folder, id, 'uv_unposed_no_shape_one_hot'), exist_ok=True)
-        # for i in range(uv_unposed_no_shape_one_hot.shape[0]):
-        #     np.savez_compressed(os.path.join(folder, id, 'uv_unposed_no_shape_one_hot', garment + '_' + str(i) + '.npz'), uv_unposed_no_shape=uv_unposed_no_shape_one_hot[i, :, :, :])
-
-        # os.makedirs(os.path.join(folder, id, 'uv_shaped_body'), exist_ok=True)
-        # np.savez_compressed(os.path.join(folder, id, 'uv_shaped_body.npz'), uv_shaped_body=uv_shaped_body)
-
-        # print(time.perf_counter() - start)
-        
-        # mesh = uv_to_mesh(uv_verts_static, m2uv)
-        # verify_reconstruction(verts_static * mask[:, np.newaxis], mesh * mask[:, np.newaxis])
-
-        # mesh = uv_to_mesh(uv_body[0,:,:], m2uv)
-        # verify_reconstruction(body[0,:,:] * mask[:, np.newaxis], mesh * mask[:, np.newaxis])
-
-        pass
-
+    h5file_curv.close()

data/curvatures.py

@@ -103,37 +103,43 @@ class Cache:
 
 
     def read(self, main_id):
-        frames = [frame for frame in self.frames if frame.split()[0] == main_id]
-        for frame in frames:
-            id = frame.split()[0]
-            garment = frame.split()[2]
-            if id in self.cache.keys():
-                continue
+        # frames = [frame for frame in self.frames if (frame.split()[0] == main_id.split()[0] and frame.split()[2] == main_id.split()[1])]
+        # for frame in frames:
+        id = main_id.split()[0]
+        garment = main_id.split()[1]
 
-            self.cache[id] = Data()
+        self.cache[f'{id} {garment}'] = Data()
 
-            mask = read_garment_mask(self.folder, id, garment)
+        is_skirt = (garment == 'Dress' or garment == 'Skirt')
 
-            _, faces = _maskFaces(np.zeros((14475,3)), F, mask)
-            v = read_posed_garment_frame(folder, id, garment, 0)[mask,:]
-            m = tm.Trimesh(v, faces, process=False, validate=False)
-            outliers = np.array([len(v) for v in m.vertex_neighbors]) == 0
-            outliers = np.nonzero(outliers)[0]
-            outliers_n_neighbours = get_closest_n_points(v, outliers, 2)
-            outliers_faces = np.concatenate((outliers[:, np.newaxis], outliers_n_neighbours), axis=-1)
-            faces = np.concatenate((faces, outliers_faces), axis=0)
-            m = tm.Trimesh(v, faces, process=False, validate=False)
+        mask = read_garment_mask(self.folder, id, garment)
         
-            if self.include_neighbours:
+        _, faces = _maskFaces(np.zeros((14475,3)), F if not is_skirt else F_skirt, mask)
+        v = read_posed_garment_frame(folder, id, garment, 0)[mask,:]
+        m = tm.Trimesh(v, faces, process=False, validate=False)
+        outliers = np.array([len(v) for v in m.vertex_neighbors]) == 0
+        outliers = np.nonzero(outliers)[0]
+        outliers_n_neighbours = get_closest_n_points(v, outliers, 2)
+        outliers_faces = np.concatenate((outliers[:, np.newaxis], outliers_n_neighbours), axis=-1)
+        # sorted_faces = [sorted(f) for f in faces.astype(np.int64)]
+        # for outlier_face in outliers_faces:
+        #     assert sorted(outlier_face) not in sorted_faces
+        faces = np.concatenate((faces, outliers_faces), axis=0)
+        m = tm.Trimesh(v, faces, process=False, validate=False)
+
+        if self.include_neighbours:
+            if not is_skirt:
                 k_neighbors = np.argsort(dist[mask, :][:, mask], axis=-1)[:, :self.neighbourhood]
-                self.cache[id].k_neighbors = k_neighbors
-                G = tm.graph.vertex_adjacency_graph(m)
-                for i in range(v.shape[0]):
-                    for k in range(25):
-                        assert nx.has_path(G, i, k_neighbors[i, k]), f'No path between {i} and {k_neighbors[i, k]}'
+            else:
+                k_neighbors = np.argsort(dist_skirt[mask, :][:, mask], axis=-1)[:, :self.neighbourhood]
+            self.cache[f'{id} {garment}'].k_neighbors = k_neighbors
+            G = tm.graph.vertex_adjacency_graph(m)
+            for i in range(v.shape[0]):
+                for k in range(25):
+                    assert nx.has_path(G, i, k_neighbors[i, k]), f'No path between {i} and {k_neighbors[i, k]}'
 
-            self.cache[id].mask = mask
-            self.cache[id].faces = faces
+        self.cache[f'{id} {garment}'].mask = mask
+        self.cache[f'{id} {garment}'].faces = faces
 
     def __getitem__(self, id):
         if id not in self.cache.keys():
@@ -164,10 +170,10 @@ def calculate_dist(folder, neighbourhood, frames):
 def generate_curvatures(folder, neighbourhood, frames):
     # ids = os.listdir(folder)[int(args.start):int(args.end)]
     # frames = frames[int(args.start):int(args.end)]
-    os.makedirs(os.path.join('/'.join(folder.split('/')[:-1]), 'backup', f'mesh_curvature_raw'), exist_ok=True)
+    os.makedirs(os.path.join(folder + '_extras', f'mesh_curvature_raw'), exist_ok=True)
 
     cache = Cache(folder, neighbourhood, frames)
-    filters = tb.Filters(complevel=9, complib='blosc2')
+    # filters = tb.Filters(complevel=9, complib='blosc2')
 
     for frame in tqdm.tqdm(frames):
 
@@ -175,29 +181,29 @@ def generate_curvatures(folder, neighbourhood, frames):
         frame_nr = int(frame.split()[1])
         garment = frame.split()[2]
     
-        mask = cache[id].mask
-        faces = cache[id].faces
+        mask = cache[f'{id} {garment}'].mask
+        faces = cache[f'{id} {garment}'].faces
         v = read_posed_garment_frame(folder, id, garment, frame_nr)[mask,:]
         
         m = tm.Trimesh(v, faces, process=False, validate=False)
         export_to_ply(m, temp_file)
         cloud = PyntCloud.from_file(temp_file)
 
-        k_neighbors = cache[id].k_neighbors
+        k_neighbors = cache[f'{id} {garment}'].k_neighbors
         ev = cloud.add_scalar_field("eigen_values", k_neighbors=k_neighbors)
         cloud.add_scalar_field("curvature", ev=ev)
         curvature = np.array(cloud.points["curvature(" + str(neighbourhood+1) + ")"])
         out = np.zeros((14475), dtype=np.float32)
-        out[mask] = curvature
+        out[mask] = curvature.real if str(curvature.dtype) == 'complex64' else curvature
 
-        np.savez_compressed(os.path.join('/'.join(folder.split('/')[:-1]), 'backup', f'mesh_curvature_raw', f'{id}_{frame_nr}_{garment}.npz'), mesh_curvature_raw=out)
-        group = h5file.get_node(f'/{id}/{frame_nr}/{garment}')
-        if f'curvature_raw' in group:
-            array = h5file.get_node(f'/{id}/{frame_nr}/{garment}/curvature_raw')
-            array[:] = out
-        else:
-            carray = h5file.create_carray(group, "curvature_raw", atom=tb.Float32Atom(), shape=out.shape, filters=filters)
-            carray[:] = out
+        np.savez_compressed(os.path.join(folder + '_extras', f'mesh_curvature_raw', f'{id}_{frame_nr}_{garment}.npz'), mesh_curvature_raw=out)
+        # group = h5file.get_node(f'/{id}/{frame_nr}/{garment}')
+        # if f'curvature_raw' in group:
+        #     array = h5file.get_node(f'/{id}/{frame_nr}/{garment}/curvature_raw')
+        #     array[:] = out
+        # else:
+        #     carray = h5file.create_carray(group, "curvature_raw", atom=tb.Float32Atom(), shape=out.shape, filters=filters)
+        #     carray[:] = out
 
     return None
 
@@ -214,34 +220,34 @@ def get_bins(values, nr_of_bins):
 def generate_curvature_labels(folder, bins, curvatures, frames, neighbourhood):
     cache = Cache(folder, neighbourhood, frames, include_neighbours=False)
 
-    os.makedirs(os.path.join('/'.join(folder.split('/')[:-1]), 'backup', f'curvature_mesh_global'), exist_ok=True)
-    os.makedirs(os.path.join('/'.join(folder.split('/')[:-1]), 'backup', f'curvature_mesh_garment'), exist_ok=True)
+    os.makedirs(os.path.join(folder + '_extras', f'curvature_mesh_global'), exist_ok=True)
+    os.makedirs(os.path.join(folder + '_extras', f'curvature_mesh_garment'), exist_ok=True)
     labels = np.linspace(0, bins.shape[0] - 2, bins.shape[0] - 1)
-    filters = tb.Filters(complevel=9, complib='blosc2')
+    # filters = tb.Filters(complevel=9, complib='blosc2')
 
     for i, frame in enumerate(tqdm.tqdm(frames)):
         id = frame.split()[0]
         frame_nr = int(frame.split()[1])
         garment = frame.split()[2]
 
-        curvature = curvatures[i][cache[id].mask]
+        curvature = curvatures[i][cache[f'{id} {garment}'].mask]
 
         frame_labels = np.zeros((curvature.shape[0]), dtype=np.uint8)
         for j in range(bins.shape[0] - 2):
             frame_labels[(curvature >= bins[j]) * (curvature < bins[j+1])] = labels[j]
         frame_labels[curvature >= bins[-2]] = labels[-1]
         out = np.zeros((14475), dtype=np.uint8)
-        out[cache[id].mask] = frame_labels + 1
+        out[cache[f'{id} {garment}'].mask] = frame_labels + 1
 
-        group = h5file.get_node(f'/{id}/{frame_nr}/{garment}')
-        if f'curvature_mesh_global' in group:
-            array = h5file.get_node(f'/{id}/{frame_nr}/{garment}/curvature_mesh_global')
-            array[:] = out
-        else:
-            carray = h5file.create_carray(group, "curvature_mesh_global", atom=tb.UInt8Atom(), shape=out.shape, filters=filters)
-            carray[:] = out
+        # group = h5file.get_node(f'/{id}/{frame_nr}/{garment}')
+        # if f'curvature_mesh_global' in group:
+        #     array = h5file.get_node(f'/{id}/{frame_nr}/{garment}/curvature_mesh_global')
+        #     array[:] = out
+        # else:
+        #     carray = h5file.create_carray(group, "curvature_mesh_global", atom=tb.UInt8Atom(), shape=out.shape, filters=filters)
+        #     carray[:] = out
 
-        np.savez_compressed(os.path.join('/'.join(folder.split('/')[:-1]), 'backup', f'curvature_mesh_global', f'{id}_{frame_nr}_{garment}.npz'), mesh_curvature_raw=out)
+        np.savez_compressed(os.path.join(folder + '_extras', f'curvature_mesh_global', f'{id}_{frame_nr}_{garment}.npz'), mesh_curvature_raw=out)
         # writePC2(os.path.join(garment_dir, "curvature_label_" + str(nr_of_labels) + "_" + str(neighbourhood) + ".pc16"), unmasked_label, True)
 
         # redo it with local bins
@@ -251,17 +257,17 @@ def generate_curvature_labels(folder, bins, curvatures, frames, neighbourhood):
             frame_labels[(curvature >= bins[j]) * (curvature < bins[j+1])] = labels[j]
         frame_labels[curvature >= bins[-2]] = labels[-1]
         out = np.zeros((14475), dtype=np.uint8)
-        out[cache[id].mask] = frame_labels + 1
+        out[cache[f'{id} {garment}'].mask] = frame_labels + 1
 
-        group = h5file.get_node(f'/{id}/{frame_nr}/{garment}')
-        if f'curvature_mesh_garment' in group:
-            array = h5file.get_node(f'/{id}/{frame_nr}/{garment}/curvature_mesh_garment')
-            array[:] = out
-        else:
-            carray = h5file.create_carray(group, "curvature_mesh_garment", atom=tb.UInt8Atom(), shape=out.shape, filters=filters)
-            carray[:] = out
+        # group = h5file.get_node(f'/{id}/{frame_nr}/{garment}')
+        # if f'curvature_mesh_garment' in group:
+        #     array = h5file.get_node(f'/{id}/{frame_nr}/{garment}/curvature_mesh_garment')
+        #     array[:] = out
+        # else:
+        #     carray = h5file.create_carray(group, "curvature_mesh_garment", atom=tb.UInt8Atom(), shape=out.shape, filters=filters)
+        #     carray[:] = out
 
-        np.savez_compressed(os.path.join('/'.join(folder.split('/')[:-1]), 'backup', f'curvature_mesh_garment', f'{id}_{frame_nr}_{garment}.npz'), mesh_curvature_raw=out)
+        np.savez_compressed(os.path.join(folder + '_extras', f'curvature_mesh_garment', f'{id}_{frame_nr}_{garment}.npz'), mesh_curvature_raw=out)
        
 
 def visualize_curvature_labels(folder, frames):
@@ -319,6 +325,7 @@ def read_raw_curvatures(folder, frames, neighbourhood):
     
     return curvatures_all
 
+
 def read_raw_curvatures_2(folder, frames, neighbourhood):
     curvatures_all = []
     for frame in tqdm.tqdm(frames):
@@ -331,17 +338,17 @@ def read_raw_curvatures_2(folder, frames, neighbourhood):
     
     return curvatures_all
 
+
 if __name__ == '__main__':
     parser = argparse.ArgumentParser()
     parser.add_argument("-f", "--folder", help = "process the given folder", default=r'/home/workspace/data/dataset/cloth3d/train')
-    parser.add_argument("-s", "--start", help = "start of range", default=0)
-    parser.add_argument("-e", "--end", help = "end of range", default=1000)
+    parser.add_argument("-s", "--start", help = "start of range", default=300)
+    parser.add_argument("-e", "--end", help = "end of range", default=600)
     parser.add_argument("-n", "--name", help = "end of range", default="lhunor")
     args = parser.parse_args()
 
-    is_skirt = False
-    skirt_modifier = '_skirt' if is_skirt else ''
-    F = np.load("custom_utils/smpl_faces" + skirt_modifier + ".npy").astype(np.int32)
+    F = np.load("custom_utils/smpl_faces.npy").astype(np.int32)
+    F_skirt = np.load("custom_utils/smpl_faces_skirt.npy").astype(np.int32)
     # k_neighbours(k=25)
 
 
@@ -352,13 +359,10 @@ if __name__ == '__main__':
     neighbourhood = 25
 
     temp_file = '/home/workspace/local-data/temp/' + ''.join(random.choice(string.ascii_lowercase) for i in range(10)) + '.ply'
-    if not is_skirt:
-        dist = np.load('custom_data/dijkstra.npy')
-    else:
-        dist = np.load('custom_data/dijkstra_skirt.npy')
+    dist = np.load('custom_data/dijkstra.npy')
+    dist_skirt = np.load('custom_data/dijkstra_skirt.npy')
 
-    print('SKIRT' if is_skirt else 'NOT SKIRT')
-    print(f'params: -f {args.folder} -s {args.start} -e {args.end} -n {args.name} -skirt {is_skirt} -neigh {neighbourhood}')
+    print(f'params: -f {args.folder} -s {args.start} -e {args.end} -n {args.name} -neigh {neighbourhood}')
 
     # frames = sorted(os.listdir(os.path.join(folder, 'mesh_posed')))[int(args.start):int(args.end)]
     # frames = [frame.split('.')[0] for frame in frames]
@@ -371,21 +375,21 @@ if __name__ == '__main__':
 
     root = '/'.join(folder.split('/')[:-1])
     type = folder.split('/')[-1]
-    h5file = tb.open_file(os.path.join(root, f'{type}_extras.h5'), mode='a')
+    # h5file = tb.open_file(os.path.join(root, f'{type}_extras.h5'), mode='a')
 
     start = time.perf_counter()
+    # frames = [f'05235 {i} Tshirt' for i in range(300)]
 
-    # generate_curvatures(folder, neighbourhood, frames)
-    curvatures = read_raw_curvatures(folder, frames, neighbourhood)
+    generate_curvatures(folder, neighbourhood, frames)
+    # curvatures = read_raw_curvatures_2(folder, frames, neighbourhood)
     # bins = get_bins(curvatures, nr_of_labels)
     # np.save(f'custom_data/{type}_bins{skirt_modifier}_{neighbourhood}.npy', bins)
     # print(bins)
-    bins = np.load(f'custom_data/{type}_bins{skirt_modifier}_{neighbourhood}.npy')
-    generate_curvature_labels(folder, bins, curvatures, frames, neighbourhood)
+    # bins = np.load(f'custom_data/{type}_bins_{neighbourhood}.npy')
+    # generate_curvature_labels(folder, bins, curvatures, frames, neighbourhood)
     # visualize_curvature_labels(folder, frames)
 
     print(time.perf_counter() - start)
 
-    h5file.close()
 
     pass

data/investigate.py

@@ -0,0 +1,160 @@
+import os
+import sys
+import random
+import numpy as np
+import plotly.graph_objects as go
+
+sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
+from data.cloth3d_io import read_garment_mask, read_garment_template, read_posed_garment_frame, loadInfo, read_posed_body_frame, read_posed_body, read_posed_garment, read_curvatures, read_curvature_frame, writeOBJ
+from custom_utils.custom_utils import _maskFaces
+
+
+def get_frames(folder, frames_file):
+    frames = []
+    with open(os.path.join(folder, frames_file), 'r') as f:
+        for line in f:
+            frames.append(line.strip())
+    return frames
+
+
+def plot(mesh, faces):
+    fig = go.Figure()
+
+    fig.add_trace(go.Mesh3d(
+                x=mesh[:,0],
+                y=mesh[:,1],
+                z=mesh[:,2],
+                i = faces[:,0],
+                j = faces[:,1],
+                k = faces[:,2]
+            ),)
+    
+    fig.update_layout(scene_aspectmode='manual', scene_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,
+                        )
+
+    fig.show()
+
+
+def build_lookup_table(repeated):
+    # iterate through the list and create a dict from each sub array in the list where each dict has one element of the sub array as key and the other elements as value as a list and repeat for each element
+
+    lookup = {}
+    for i in range(len(repeated)):
+        for j in range(len(repeated[i])):
+            if len(repeated[i]) == 1:
+                continue
+            lookup[repeated[i][j]] = [repeated[i][k] for k in range(len(repeated[i]))] # if k != j
+    return lookup
+
+def mask_faces(all_faces, mask, repeated, verts):
+    # if any of the vertices of a face is masked, mask the face
+    lookup = build_lookup_table(repeated)
+    faces = []
+    for face in all_faces:
+        # if any two or all three vertices of the face have the same coordinates then continue
+        if np.allclose(verts[face[0]], verts[face[1]]) or np.allclose(verts[face[0]], verts[face[2]]) or np.allclose(verts[face[1]], verts[face[2]]):
+            continue
+
+        if mask[face[0]] and mask[face[1]] and mask[face[2]]:
+            faces.append(face)
+        else:
+            # create three arrays of vertices from the lookup table for each vertex of the face
+            # but only for those vertices that are in the lookup table, otherwise createa a single element array with the vertex
+            v1 = lookup[face[0]] if face[0] in lookup else [face[0]]
+            v2 = lookup[face[1]] if face[1] in lookup else [face[1]]
+            v3 = lookup[face[2]] if face[2] in lookup else [face[2]]
+            # try each combination of vertices from the three arrays and check if the face is masked
+            found_replacement = False
+            for i in range(len(v1)):
+                if not found_replacement:
+                    for j in range(len(v2)):
+                        if not found_replacement:
+                            for k in range(len(v3)):
+                                if not found_replacement:
+                                    if mask[v1[i]] and mask[v2[j]] and mask[v3[k]]:
+                                        # if the euclidean distance between any two of the vertices of the face is less than 0.05 then append the face
+                                        # if np.linalg.norm(verts[v1[i]] - verts[v2[j]]) < 0.05 and np.linalg.norm(verts[v1[i]] - verts[v3[k]]) < 0.05 and np.linalg.norm(verts[v2[j]] - verts[v3[k]]) < 0.05:
+                                        faces.append(np.array([v1[i], v2[j], v3[k]], dtype=np.int32))
+                                        assert np.all(verts[face] == verts[faces[-1]])
+                                        found_replacement = True
+                                        break
+                # for j in range(len(v2)):
+                #     for k in range(len(v3)):
+                #         if mask[v1[i]] and mask[v2[j]] and mask[v3[k]]:
+                #             faces.append([v1[i], v2[j], v3[k]])
+                #             break
+
+    faces = np.array(faces)
+    V_ind = np.zeros(mask.shape, dtype=np.int32)
+    V_ind[mask] = np.arange(mask.sum())
+    F_new = np.take(V_ind, faces.reshape(-1), axis=0).reshape((-1, 3))
+    return F_new
+
+
+def check(folder, frame):
+    id, frame_nr, garment = frame.split()
+    frame_nr = int(frame_nr)
+    info = loadInfo(os.path.join(folder, id, 'info.mat'))
+    mask = read_garment_mask(folder, id, garment)
+    mesh_posed = read_posed_garment_frame(folder, id, garment, frame_nr, mask=mask, shape=info['shape'], pose=info['poses'][:,frame_nr], deorient=use_no_orient, unshape=no_shape)
+
+    v = mesh_posed[mask,:]
+    is_skirt = (garment == 'Dress' or garment == 'Skirt')
+    skirt_mod = '_skirt' if is_skirt else ''
+    F = np.load(f'custom_utils/smpl_faces{skirt_mod}.npy').astype(np.int32)
+    _, faces = _maskFaces(np.zeros((14475,3)), F, mask)
+
+    # plot(v, faces)
+
+    unq, idx, count = np.unique(mesh_posed, axis=0, return_counts=True, return_index=True)
+    repeated_groups = unq[count > 1]
+
+    mask_2 = np.zeros((14475), dtype=bool)
+    mask_2[idx] = True
+    v_2 = mesh_posed[mask_2,:]
+    _, faces_2 = _maskFaces(np.zeros((14475,3)), F, mask_2)
+    # plot(v_2, faces_2)
+
+    # writeOBJ('temp/investigate.obj', v, faces)
+
+    repeated = []
+    for repeated_group in repeated_groups:
+        repeated_idx = np.argwhere(np.all(mesh_posed == repeated_group, axis=1))
+        repeated.append(repeated_idx.ravel())
+
+    faces_3 = mask_faces(F, mask_2, repeated, mesh_posed)
+    writeOBJ('temp/investigate_3.obj', v_2, faces_3)
+    plot(v_2, faces_3)
+
+
+    pass
+    
+
+
+if __name__ == "__main__":
+    folder = '/home/workspace/data/dataset/cloth3d/train'
+    frames = get_frames(folder, 'frames_all.txt')
+    root = '/'.join(folder.split('/')[:-1])
+    type = folder.split('/')[-1]
+
+    use_no_orient = False
+    no_shape = False
+    is_skirt = True
+
+    skirt_modifier = '_skirt' if is_skirt else ''
+    orient_mod = '_no_orient' if use_no_orient else ''
+    shape_mod = '_no_shape' if no_shape else ''
+
+
+    while True:
+        frame = random.choice(frames)
+        frame = '05936 0 Dress'
+
+        check(folder, frame)
+
+        pass

data/preprocessing.py

@@ -111,7 +111,7 @@ def create_extras_file(folder):
     # create a new pytables file
     root = '/'.join(folder.split('/')[:-1])
     type = folder.split('/')[-1]
-    h5file = tb.open_file(os.path.join(root, f'{type}_extras.h5'), mode='w', title='Extras')
+    h5file = tb.open_file(os.path.join(root, f'{type}_extras_test.h5'), mode='w', title='Extras')
 
     frames = []
     with open(os.path.join(folder, 'frames_all.txt'), 'r') as f:
@@ -119,7 +119,7 @@ def create_extras_file(folder):
         lines = [line.strip() for line in lines]
     frames = lines
 
-    for frame in tqdm(frames[:3000]):
+    for frame in tqdm(frames[:1000]):
         # add a new empty group to the root group for each frame
         id, frame_nr, garment = frame.split()
         if not id in h5file.root:

data/verify.py

@@ -0,0 +1,138 @@
+import os
+import sys
+import random
+import numpy as np
+import tables as tb
+import plotly.graph_objects as go
+import plotly.express as ex
+from plotly.subplots import make_subplots
+
+sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
+from data.cloth3d_io import read_garment_mask, read_garment_template, read_posed_garment_frame, loadInfo, read_posed_body_frame, read_posed_body, read_posed_garment, read_curvatures, read_curvature_frame
+from custom_utils.custom_utils import _maskFaces
+
+
+def get_frames(folder, frames_file):
+    frames = []
+    with open(os.path.join(folder, frames_file), 'r') as f:
+        for line in f:
+            frames.append(line.strip())
+    return frames
+
+
+def verify_uv(folder, frame):
+    id, frame_nr, garment = frame.split()
+    uv_static = np.load(os.path.join(folder, 'uv_static', f'{id}_{garment}.npz'))['uv_static']
+    uv_mask = np.load(os.path.join(folder, 'uv_mask', f'{id}_{garment}.npz'))['uv_mask']
+    uv_posed = np.load(os.path.join(folder, 'uv_posed' + orient_mod, f'{id}_{frame_nr}_{garment}.npz'))['uv_posed']
+    uv_body_posed = np.load(os.path.join(folder, 'uv_body_posed' + orient_mod, f'{id}_{frame_nr}_{garment}.npz'))['uv_body_posed']
+    uv_curvature = np.load(os.path.join(folder, 'uv_curvature', f'{id}_{frame_nr}_{garment}.npz'))['uv_curvature']
+
+    all_1 = np.concatenate([uv_static, uv_posed, uv_body_posed], axis=1)
+    all_1 = (all_1 - all_1.min()) / (all_1.max() - all_1.min())
+
+    ex.imshow(all_1).show()
+    ex.imshow(uv_mask).show()
+    ex.imshow(uv_curvature).show()
+
+
+def verify_mesh(folder, frame):
+    id, frame_nr, garment = frame.split()
+    frame_nr = int(frame_nr)
+    info = loadInfo(os.path.join(folder, id, 'info.mat'))
+    mask = read_garment_mask(folder, id, garment)
+    mesh_static = read_garment_template(folder, id, garment, mask)[mask,:]
+    mesh_posed = read_posed_garment_frame(folder, id, garment, frame_nr, mask=mask, shape=info['shape'], pose=info['poses'][:,frame_nr], deorient=use_no_orient, unshape=no_shape)[mask,:]
+    body_posed = read_posed_body_frame(folder, id, frame_nr, info['poses'], use_no_orient)
+    mesh_curvature = read_curvature_frame(folder + '_extras', id, garment, frame_nr)[mask]
+
+    fig = make_subplots(rows=1, cols=4, start_cell="top-left", specs = [ [{ 'type': 'scene' } for _ in range(4)] ])
+    F = np.load(f'custom_utils/smpl_faces{skirt_modifier}.npy').astype(np.int32)
+    _, faces = _maskFaces(np.zeros((14475,3)), F, mask)
+
+
+    fig.add_trace(go.Mesh3d(
+                x=mesh_static[:,0],
+                y=mesh_static[:,1],
+                z=mesh_static[:,2],
+                i = faces[:,0],
+                j = faces[:,1],
+                k = faces[:,2]
+            ),
+            row=1, col=1)
+    
+    fig.add_trace(go.Mesh3d(
+                x=mesh_posed[:,0],
+                y=mesh_posed[:,1],
+                z=mesh_posed[:,2],
+                i = faces[:,0],
+                j = faces[:,1],
+                k = faces[:,2]
+            ),
+            row=1, col=2)
+
+    fig.add_trace(go.Mesh3d(
+                x=body_posed[:,0],
+                y=body_posed[:,1],
+                z=body_posed[:,2],
+                i = F[:,0],
+                j = F[:,1],
+                k = F[:,2],
+            ),
+            row=1, col=3)
+
+    fig.add_trace(go.Mesh3d(
+                x=mesh_posed[:,0],
+                y=mesh_posed[:,1],
+                z=mesh_posed[:,2],
+                i = faces[:,0],
+                j = faces[:,1],
+                k = faces[:,2],
+                vertexcolor=np.tile(mesh_curvature[:,None], (1,3)) * 40,
+            ),
+            row=1, col=4)
+
+    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),
+                    scene4_aspectmode='manual', scene4_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,
+                        scene4 = helper,
+                        )
+
+    fig.show()
+
+    pass
+
+
+if __name__ == '__main__':
+    folder = '/home/workspace/data/dataset/cloth3d/train'
+    frames = get_frames(folder, 'frames_all.txt')
+    root = '/'.join(folder.split('/')[:-1])
+    type = folder.split('/')[-1]
+    # h5file_curv = tb.open_file(os.path.join(root, f'{type}_extras.h5'), mode='r')
+
+    use_no_orient = False
+    no_shape = False
+    is_skirt = True
+
+    skirt_modifier = '_skirt' if is_skirt else ''
+    orient_mod = '_no_orient' if use_no_orient else ''
+    shape_mod = '_no_shape' if no_shape else ''
+
+
+    while True:
+        frame = random.choice(frames)
+        frame = '05936 0 Dress'
+
+        verify_mesh(folder, frame)
+        verify_uv(folder + '_extras', frame)
+
+        pass
+

run_parallel.py

@@ -0,0 +1,29 @@
+import os
+
+
+
+if __name__ == '__main__':
+
+    folder = '/home/workspace/data/dataset/cloth3d/train'
+    # folder = '/home/workspace/local-data/cape_processed/train/male'
+    # folder = '/home/workspace/local-data/valid_skirt'
+    # folder = '/home/workspace/local-data/valid'
+    # ids = os.listdir(folder)
+    # ids = os.listdir(folder + '/mesh_posed_no_orient')#[280000:]
+    skirt = 0
+
+    count = int(500 * 1e3)
+    step = 10000
+
+    # !!!!!!! change skirt parameter in script
+
+    while count < 510 * 1e3:
+        # os.system('nohup python convert_to_uv_cape.py -f ' + folder + ' -s ' + str(count) + ' -e ' + str(count + step) + ' -n ' + 'to_uv_2_cape_train_male_' + str(count) + '_' + str(count+step) + ' > nohup/to_uv_2_cape_train_male_' + str(count) + '_' + str(count+step) + '.out' + ' 2>&1 &')
+        # os.system('nohup python convert_to_uv.py -f ' + folder + ' -s ' + str(count) + ' -e ' + str(count + step) + ' -n ' + 'to_uv_train_unposed_onehot_' + str(count) + '_' + str(count+step) + ' > nohup/to_uv_train_unposed_onehot_' + str(count) + '_' + str(count+step) + '.out' + ' 2>&1 &')
+        # os.system('nohup python convert_to_uv_local.py -f ' + folder + ' -s ' + str(count) + ' -e ' + str(count + step) + ' -n ' + 'to_uv_train_local_unposed_onehot_' + str(count) + '_' + str(count+step) + ' > nohup/to_uv_train_local_unposed_onehot_' + str(count) + '_' + str(count+step) + '.out' + ' 2>&1 &')
+        os.system('nohup python data/convert_to_uv.py -f ' + folder + ' -s ' + str(count) + ' -e ' + str(count + step) + ' -n ' + 'to_uv_train_2_' + str(count) + '_' + str(count+step) + ' > nohup/to_uv/train_2_' + str(count) + '_' + str(count+step) + '_again.out' + ' 2>&1 &')
+        # os.system('nohup python convert_to_uv_local_single.py -f ' + folder + ' -s ' + str(count) + ' -e ' + str(count + step) + ' -n ' + 'to_uv_train_skirt_local_unposed_onehot_3_' + str(count) + '_' + str(count+step) + ' > nohup/to_uv_train_skirt_local_unposed_onehot_3_' + str(count) + '_' + str(count+step) + '.out' + ' 2>&1 &')
+        # os.system('nohup python convert_to_uv_local_single.py -f ' + folder + ' -s ' + str(count) + ' -e ' + str(count + step) + ' -n ' + 'to_uv_valid_skirt_local_unposed_onehot_3_' + str(count) + '_' + str(count+step) + ' > nohup/to_uv_valid_skirt_local_unposed_onehot_' + str(count) + '_' + str(count+step) + '.out' + ' 2>&1 &')
+        # os.system('nohup python downscale.py -f ' + folder + ' -s ' + str(count) + ' -e ' + str(count + step) + ' -n ' + 'downscale_valid_posed_no_orient_' + str(count) + '_' + str(count+step) + ' > nohup/downscale_valid_posed_no_orient_' + str(count) + '_' + str(count+step) + '.out' + ' 2>&1 &')
+        # os.system('python Data/convert_to_uv.py -f ' + folder + ' -s ' + str(count) + ' -e ' + str(count + step) + '')
+        count += step

scripts/cvae.py

@@ -26,37 +26,25 @@ import numpy as np
 import plotly.graph_objects as go
 import plotly.express as ex
 from PIL import Image
-from postprocessing import fix_collision
+from utils.postprocess import fix_collision
 
 if config['is_skirt']:
-    from losses import edge_loss_skirt as edge_loss
+    from utils.losses import edge_loss_skirt as edge_loss
 else:
-    from losses import edge_loss
+    from utils.losses import edge_loss
 
-
-
-from losses import collision_ratio, kl_loss, normal_loss #, concept_loss
-# from losses import collision_penalty as collision_penalty
-from losses import collision_loss as collision_penalty
+from utils.losses import collision_ratio, kl_loss, normal_loss, collision_loss
 from custom_utils.custom_utils import batch_to_dict, load_model, plot_3d, UV2Mesh, get_summary_html, maskFacesBatch, _maskFaces, plot_3d_collision
-from dataset import DatasetFactory, StaticDataset
+from data.dataset import DatasetFactory, StaticDataset
 from custom_utils.smpl import SMPL
 from models.cvae_model import CVAENew as CVAE
+from utils.utils import Accumulator
 
 
 def init_wandb():
     run = wandb.init(project='ablation', name=config['name'], notes=config['notes'], dir='/wandb', tags=config['tags'])
     artifact = wandb.Artifact('code', type='code')
-    artifact.add_file('cvae.py')
-    artifact.add_file('models/cvae_model.py')
-    artifact.add_file('losses.py')
-    # artifact.add_file('configs/config_base.py')
-    # artifact.add_file('configs/config_cvae.py')
-    artifact.add_file('models/base.py')
-    artifact.add_file('models/cond_encoder.py')
-    artifact.add_file('dataset.py')
-    artifact.add_file('convnext/models/convnext_tf.py')
-    config_file_names = [key for key in sys.modules.keys()  if key.startswith('configs.')]
+    config_file_names = [key for key in sys.modules.keys()  if (key.startswith('configs.') or key.startswith('custom_utils.') or key.startswith('utils.') or key.startswith('data.') or key.startswith('models.') or key.startswith('scripts.'))]
     for config_file_name in config_file_names:
         file_name = '/'.join(config_file_name.split('.')) + '.py'
         if os.path.exists(file_name):
@@ -101,7 +89,7 @@ def train_step(batch):
         # is the output already masked?
 
         e_loss = tf.reduce_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 = collision_loss(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask'])
         c_loss, c_ratio = tf.reduce_mean(c_loss), tf.reduce_mean(c_ratio)
         l2_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(((uv_orig - pred) ** 2), axis=-1), axis=[1, 2]))
         l2_mesh_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (batch['mesh_posed'] - mesh_pred_posed) ** 2, axis=-1), axis=-1))
@@ -159,7 +147,7 @@ def valid_step(batch):
     # is the output already masked?
 
     e_loss = tf.reduce_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 = collision_loss(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask'])
     c_loss, c_ratio = tf.reduce_mean(c_loss), tf.reduce_mean(c_ratio)
     l2_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(((uv_orig - pred) ** 2), axis=-1), axis=[1, 2]))
     l2_mesh_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (batch['mesh_posed'] - mesh_pred_posed) ** 2, axis=-1), axis=-1))
@@ -219,7 +207,7 @@ def test_step(batch):
     # is the output already masked?
 
     e_loss = tf.reduce_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 = collision_loss(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask'])
     c_loss, c_ratio = tf.reduce_mean(c_loss), tf.reduce_mean(c_ratio)
     l2_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(((uv_orig - pred) ** 2), axis=-1), axis=[1, 2]))
     l2_mesh_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1), axis=-1))
@@ -426,9 +414,6 @@ def sample(folder):
         print()
 
 
-
-
-
 def test():
     str_split = config['name'].split('_')
     test_idx = str_split.index('test')
@@ -518,6 +503,9 @@ if __name__ == '__main__':
                                          in_memory=True, 
                                         is_cape=config['is_cape'])
 
+    train_accumulator = Accumulator()
+    valid_accumulator = Accumulator()
+
     model = CVAE(config)
 
     optimizer = tf.optimizers.Adam(learning_rate=config['lr'])
@@ -530,7 +518,6 @@ if __name__ == '__main__':
     ckpt = tf.train.Checkpoint(model=model)
     manager = tf.train.CheckpointManager(ckpt, 'checkpoints/post/' + config['name'], max_to_keep=20)
 
-
     if config['pretrained_model'] != '':
         load_model(model, config['pretrained_model'])
 
@@ -542,30 +529,39 @@ if __name__ == '__main__':
         batch = static_dataset_train.expand_batch(batch_to_dict(batch, config['signature_dynamic']))
 
         train_loss = train_step(batch)
+        train_accumulator.update(train_loss)
 
         data_valid = valid_iterator.get_next()
         data_valid = static_dataset_valid.expand_batch(batch_to_dict(data_valid, config['signature_dynamic']))
         valid_loss = valid_step(data_valid)
+        valid_accumulator.update(valid_loss)
 
-        to_log = { **train_loss, **valid_loss, **{ 'current_step': step} }
-        if step % 5000 == 0 or (step < 5000 and step % 500 == 0):
-            train_figs = plot_everything(batch, 'train')
-            valid_figs = plot_everything(data_valid, 'valid')
-            to_log = { **to_log, **train_figs, **valid_figs }
 
-        if step % 5000 == 0:
-            manager.save(checkpoint_number=step)
-        if step == 0:
-            to_log = { **to_log, **{ 'model/summary': wandb.Html(get_summary_html(model)) } }
-
-        if train_loss['train/euclid_error'] < (config['balancing_activation_limit'] if 'balancing_activation_limit' in config.keys() else 0.017):
+        if train_accumulator.get()['train/euclid_error'] < (config['balancing_activation_limit'] if 'balancing_activation_limit' in config.keys() else 0.017):
             if balancing_activated_step == 0:
                 balancing_activated_step = step
             g_balancing_coeff = g_balancing_coeff + 1 / 10000     
 
-        to_log = { **to_log, **{ 'train/balancing_coeff': g_balancing_coeff } }
 
-        wandb.log(to_log, step = step)
+        if step % 1000 == 0 or (step < 1000 and step % 100 == 0):
+            to_log = { **train_accumulator.get(), **valid_accumulator.get(), **{ 'current_step': step} }
+
+            if step % 5000 == 0 or (step < 5000 and step % 1000 == 0):
+                train_figs = plot_everything(batch, 'train')
+                valid_figs = plot_everything(data_valid, 'valid')
+                to_log = { **to_log, **train_figs, **valid_figs }
+
+            if step == 0:
+                to_log = { **to_log, **{ 'model/summary': wandb.Html(get_summary_html(model)) } }
+
+            to_log = { **to_log, **{ 'train/balancing_coeff': g_balancing_coeff } }
+        
+            wandb.log(to_log, step = step)
+            train_accumulator.reset()
+            valid_accumulator.reset()
+            
+        if step % 5000 == 0:
+            manager.save(checkpoint_number=step)
 
         print(step)
 

utils/losses.py

@@ -0,0 +1,270 @@
+import tensorflow as tf
+
+import numpy as np
+import tensorflow_graphics as tfg
+import scipy as sp
+import scipy.io as sio
+
+from scipy.spatial import cKDTree
+from tensorflow_graphics.geometry.representation.mesh.normals import vertex_normals
+from tensorflow_graphics.geometry.convolution.graph_convolution import edge_convolution_template
+
+
+
+# from hbertiche/PBNS
+""" COLLISION FUNCTIONS START """
+F_coll = np.load("custom_utils/smpl_faces.npy")
+F_coll_skirt = np.load("custom_utils/smpl_faces_skirt.npy")
+F_coll_orig = np.load("custom_utils/smpl_faces_orig.npy")
+F_norm = np.tile(F_coll[np.newaxis,:,:], [4,1,1])
+F_norm_skirt = np.tile(F_coll_skirt[np.newaxis,:,:], [4,1,1])
+F_norm_orig = np.tile(F_coll_orig[np.newaxis,:,:], [4,1,1])
+
+def _nearest_neighbour(V, B):
+	tree = cKDTree(B.numpy())
+	return tree.query(V.numpy(), workers=-1)[1]
+
+@tf.function
+def _tf_nn_parallel(V, B):
+	return tf.map_fn(fn=lambda elem: tf.py_function(func=_nearest_neighbour, inp=elem, Tout=tf.int32), elems=(V,B), fn_output_signature=tf.int32, parallel_iterations=V.shape[0])
+
+# @tf.function
+def collision_loss(V, B, mask, thr=.004):
+    B_vn = vertex_normals(B, tf.tile(F_coll[None], [B.shape[0], 1, 1]))
+   
+    # compute correspondences
+    idx = _tf_nn_parallel(V, B)
+    N, M = V.shape[:2]
+    # _idx = tf.tile(np.array(list(range(N)), np.int32)[:,None], [1, M])
+    _idx = tf.ones((N, M), dtype=tf.int32) * np.array(list(range(N)), np.int32)[:,None]
+    idx = tf.stack((_idx, idx), axis=-1)
+    # compute loss
+    D = V - tf.gather_nd(B, idx)
+    b_vn = tf.gather_nd(B_vn, idx)
+    dot = tf.einsum('abc,abc->ab', D, b_vn)
+    thresholded = tf.minimum(dot - thr, 0)  * mask
+    loss = tf.reduce_sum(thresholded ** 2, axis=-1)
+    ratio = tf.reduce_sum(tf.cast(thresholded < 0, dtype=tf.float32)) / tf.reduce_sum(mask)
+
+    return loss, ratio
+
+
+def collision_ratio(V, B, mask, thr=.004):
+    B_vn = vertex_normals(B, tf.tile(F[None], [B.shape[0], 1, 1]))
+   
+    # compute correspondences
+    idx = _tf_nn_parallel(V, B)
+    N, M = V.shape[:2]
+    # _idx = tf.tile(np.array(list(range(N)), np.int32)[:,None], [1, M])
+    _idx = tf.ones((N, M), dtype=tf.int32) * np.array(list(range(N)), np.int32)[:,None]
+    idx = tf.stack((_idx, idx), axis=-1)
+    # compute loss
+    D = V - tf.gather_nd(B, idx)
+    b_vn = tf.gather_nd(B_vn, idx)
+    dot = tf.einsum('abc,abc->ab', D, b_vn)
+    loss = tf.minimum(dot - thr, 0)  * mask # / tf.reduce_sum(mask, axis=-1)
+    ratio = tf.reduce_sum(tf.cast(loss < 0, dtype=tf.float32)) / tf.reduce_sum(mask)
+
+    return ratio
+
+""" COLLISION FUNCTIONS END """
+
+
+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)
+
+def edge_loss(x, y, mask):
+    e_mask = get_edge_mask(mask)
+    x_e = tf.gather(x, edges[:,0], axis=1) - tf.gather(x, edges[:,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[:,0], axis=1) - tf.gather(y, edges[:,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
+
+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):
+        super(PairwiseDistance, self).__init__(**kwargs)
+
+    def call(self, A, B):
+        rA = tf.reduce_sum(tf.square(A), axis=-1)
+        rB = tf.reduce_sum(tf.square(B), axis=-1)
+        transpose_axes = [0, 2, 1] 
+        distances = - 2*tf.matmul(A, tf.transpose(B, transpose_axes)) + rA[:, :, tf.newaxis] + rB[:, tf.newaxis, :]
+        return distances
+
+
+class NearestNeighbour(tf.keras.layers.Layer):
+    def __init__(self, **kwargs):
+        super(NearestNeighbour, self).__init__(**kwargs)
+
+    def call(self, A, B):
+        distances = PairwiseDistance(dtype=self.dtype)(A, B)
+        nearest_neighbour = tf.argmin(distances, axis=-1)
+        return tf.cast(nearest_neighbour, dtype=tf.int32)
+
+#
+# @tf.function
+def collision_penalty(va, vb, mask, eps=1e-3):
+    nb = vertex_normals(vb, tf.tile(F[None], [vb.shape[0], 1, 1]))
+    batch_size = tf.cast(tf.shape(va)[0], va.dtype)
+
+    closest_vertices = NearestNeighbour(dtype=va.dtype)(va, vb)
+    vb = tf.gather(vb, closest_vertices, batch_dims=1)
+    nb = tf.gather(nb, closest_vertices, batch_dims=1)
+
+    distance = tf.reduce_sum(nb*(va - vb), axis=-1) 
+    interpenetration = tf.maximum(eps - distance * mask, 0)
+
+    return tf.reduce_sum(interpenetration**3) / batch_size
+
+whole_garment_uv_mask = np.load('custom_utils/mask_256x512.npy')
+
+def uv_mask_loss(pred, uv_mask):
+    inv_mask = (1 - uv_mask)  * whole_garment_uv_mask
+    pred = tf.reduce_sum(pred, axis=1)
+    invalid = pred * inv_mask
+    invalid = invalid > 0.01
+    invalid_count = tf.math.count_nonzero(invalid, dtype=tf.float32)
+    total_count = tf.math.count_nonzero(inv_mask, dtype=tf.float32)
+    return invalid_count / total_count
+
+
+
+
+def merge_verts(Vs, batchDim):
+    V = tf.concat(tf.unstack(Vs, num=batchDim), axis=0)
+    return V
+    
+
+def _bend_loss_single(x, F, L, mask):
+    verts = tf.boolean_mask(x, mask, axis=0)
+    VN = vertex_normals(verts, F)
+    bend = edge_convolution_template(
+        VN, L, sizes=None,
+        edge_function=lambda x,y: (x - y)**2,
+        reduction='weighted',
+        edge_function_kwargs={}
+    )
+    loss = tf.reduce_sum(bend)
+    return loss
+
+def bend_loss(x, F, L, mask):
+    # x: predicted outfits
+    # F: faces
+    # L: laplacian
+    loss = tf.stack([_bend_loss_single(x[i], F[i], L[i], mask[i]) for i in range(x.shape[0])])
+    return loss
+
+
+# def bend_loss(x, F, L, mask):
+#     # x: predicted outfits
+#     # F: faces
+#     # L: laplacian
+#     # x = x * mask[:,:,tf.newaxis]
+#     batch_size = x.shape[0]
+#     verts, faces = maskFacesBatch(x, F, mask)
+#     nr_valid_verts = [tf.reduce_sum(mask[i]) for i in range(batch_size)]
+#     x = tf.concat(verts, axis=0)
+#     F = tf.concat([faces[i] + tf.reduce_sum(nr_valid_verts[:i]) for i in range(1, batch_size)])
+#     F = tf.concat((faces[0], F), axis=0)
+#     VN = vertex_normals(x, F)
+#     bend = edge_convolution_template(
+#         VN, L, sizes=None,
+#         edge_function=lambda x,y: (x - y)**2,
+#         reduction='weighted',
+#         edge_function_kwargs={}
+#     )
+#     loss = tf.reduce_sum(bend, axis=[1])
+#     loss = tf.reshape(loss, (batch_size, 14475))
+#     loss = tf.reduce_mean(loss, axis=1)
+#     return loss
+
+
+def log_normal_pdf(sample, mean, logvar, raxis=1):
+    log2pi = tf.math.log(2. * np.pi)
+    res = tf.reduce_sum(-.5 * ((sample - mean) ** 2. * tf.exp(-logvar) + logvar + log2pi), axis=raxis)
+    return res
+
+
+def kl_loss(mean, logvar):
+    kl_loss = 1 + logvar - tf.square(mean) - tf.exp(logvar)
+    kl_loss = tf.reduce_mean(kl_loss)
+    kl_loss *= -0.5
+    return kl_loss
+
+
+def normal_loss(pred, gt, mask, is_skirt=False, is_smpl_orig=False):
+    f_norm = None
+    if not is_skirt and not is_smpl_orig:
+        f_norm = F_norm
+    elif is_skirt:
+        f_norm = F_norm_skirt
+    elif is_smpl_orig:
+        f_norm = F_norm_orig
+    else:
+        raise NotImplementedError
+
+    pred_norm = vertex_normals(pred, f_norm)
+    gt_norm = vertex_normals(gt, f_norm)
+    return tf.reduce_mean(tf.reduce_mean((tf.reduce_sum(tf.math.abs(pred_norm - gt_norm), axis=-1) * mask), axis=-1))
+
+# vgg = tf.keras.applications.vgg16.VGG16(include_top=False, input_shape=(512,256,3))
+# vgg_block_1 = tf.keras.Sequential(vgg.layers[:4])
+# vgg_block_2 = tf.keras.Sequential(vgg.layers[4:7])
+# vgg_block_3 = tf.keras.Sequential(vgg.layers[7:11])
+
+# def concept_loss(pred, gt):
+#     pred = tf.keras.applications.vgg16.preprocess_input((pred + 1) * 128)
+#     gt = tf.keras.applications.vgg16.preprocess_input((gt + 1) * 128)
+
+#     pred_features_1 = vgg_block_1(pred, training=False)
+#     pred_features_2 = vgg_block_2(pred_features_1, training=False)
+#     pred_features_3 = vgg_block_3(pred_features_2, training=False)
+#     gt_features_1 = vgg_block_1(gt, training=False)
+#     gt_features_2 = vgg_block_2(gt_features_1, training=False)
+#     gt_features_3 = vgg_block_3(gt_features_2, training=False)
+
+#     l1_1 = tf.reduce_mean(tf.math.abs(pred_features_1 - gt_features_1))
+#     l1_2 = tf.reduce_mean(tf.math.abs(pred_features_2 - gt_features_2))
+#     l1_3 = tf.reduce_mean(tf.math.abs(pred_features_3 - gt_features_3))
+
+#     return l1_1 + l1_2 + l1_3
+
+
+if __name__ == '__main__':
+
+
+    V = tf.ones((4, 512, 256, 3))
+    B = tf.ones((4, 512, 256, 3))
+    m = tf.ones((4, 14475))
+    loss_1 = concept_loss(V, B)
+    print()

utils/postprocess.py

@@ -0,0 +1,137 @@
+import os
+
+# os.environ["CUDA_VISIBLE_DEVICES"] = "3"
+import tensorflow as tf
+
+import sys
+import shutil
+import tqdm
+import random
+import math
+import numpy as np
+import scipy as sp
+import scipy.io as sio
+import plotly.graph_objects as go
+import plotly.express as ex
+
+from custom_utils.custom_utils import _maskFaces
+
+from scipy.spatial import cKDTree
+from tensorflow_graphics.geometry.representation.mesh.normals import vertex_normals
+
+
+# from hbertiche/PBNS
+""" COLLISION FUNCTIONS START """
+F = np.load("custom_utils/smpl_faces.npy")
+
+
+def _nearest_neighbour(V, B):
+	tree = cKDTree(B.numpy())
+	return tree.query(V.numpy(), workers=-1)[1]
+
+@tf.function
+def _tf_nn_parallel(V, B):
+	return tf.map_fn(fn=lambda elem: tf.py_function(func=_nearest_neighbour, inp=elem, Tout=tf.int32), elems=(V,B), fn_output_signature=tf.int32, parallel_iterations=V.shape[0])
+
+
+def fix_collision(V, B, mask, thr=.004):
+    B_vn = vertex_normals(B, tf.tile(F[None], [B.shape[0], 1, 1]))
+   
+    # compute correspondences
+    idx = _tf_nn_parallel(V, B)
+    N, M = V.shape[:2]
+    # _idx = tf.tile(np.array(list(range(N)), np.int32)[:,None], [1, M])
+    _idx = tf.ones((N, M), dtype=tf.int32) * np.array(list(range(N)), np.int32)[:,None]
+    idx = tf.stack((_idx, idx), axis=-1)
+    # compute loss
+    D = V - tf.gather_nd(B, idx)
+    b_vn = tf.gather_nd(B_vn, idx)
+    dot = tf.einsum('abc,abc->ab', D, b_vn)
+    loss = tf.minimum(dot - thr, 0)  * mask
+    B_offset_closest = tf.gather_nd(B, idx) + (b_vn / 2e2)
+    output = tf.where(loss[:,:,tf.newaxis] < 0, B_offset_closest, V)
+
+    return output
+
+
+""" COLLISION FUNCTIONS END """
+
+
+
+if __name__ == '__main__':
+
+    F = np.load("custom_utils/smpl_faces.npy")
+    posed = np.load('/home/workspace/local-data/train/mesh_posed/00016_10.npz')['mesh_posed'][np.newaxis,:,:]
+    body = np.load('/home/workspace/local-data/train/mesh_body_posed/00016_10.npy')[np.newaxis,:,:]
+    mask = np.load('/home/workspace/local-data/train/mesh_mask/00016.npy')
+    posed_2 = np.load('/home/workspace/local-data/train/mesh_posed/00016_15.npz')['mesh_posed'][np.newaxis,:,:]
+
+    fixed = fix_collision(posed_2, body, mask).numpy()
+
+    x = posed_2[0, mask,:]
+    y = body[0]
+    _, faces = _maskFaces(np.zeros((14475,3), np.float32), F, mask)
+    fig = go.Figure(data=[
+            go.Mesh3d(
+                x=x[:,0],
+                y=x[:,1],
+                z=x[:,2],
+                i = faces[:,0],
+                j = faces[:,1],
+                k = faces[:,2],
+            ),
+            go.Mesh3d(
+                x=y[:,0],
+                y=y[:,1],
+                z=y[:,2],
+                i = F[:,0],
+                j = F[:,1],
+                k = F[:,2],
+            ),
+        ])
+
+    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()
+
+    x = fixed[0, mask,:]
+    y = body[0]
+    _, faces = _maskFaces(np.zeros((14475,3), np.float32), F, mask)
+    fig = go.Figure(data=[
+            go.Mesh3d(
+                x=x[:,0],
+                y=x[:,1],
+                z=x[:,2],
+                i = faces[:,0],
+                j = faces[:,1],
+                k = faces[:,2],
+            ),
+            go.Mesh3d(
+                x=y[:,0],
+                y=y[:,1],
+                z=y[:,2],
+                i = F[:,0],
+                j = F[:,1],
+                k = F[:,2],
+            ),
+        ])
+
+    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()

utils/utils.py

@@ -0,0 +1,27 @@
+import os
+import sys
+import numpy as np
+
+
+
+class Accumulator():
+    def __init__(self):
+        self.dict = dict()
+
+    def update(self, dict):
+        for key in dict.keys():
+            if key in self.dict.keys():
+                self.dict[key].append(dict[key])
+            else:
+                self.dict[key] = [dict[key]]
+
+    def get(self):
+        result = dict()
+        for key in self.dict.keys():
+            result[key] = np.mean(self.dict[key])
+        return result
+
+    def reset(self):
+        self.dict = dict()
+
+