save progress

.gitignore

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

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

@@ -244,14 +244,13 @@ 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'))
+    # 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])])
+    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, :]
@@ -298,7 +297,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,8 +308,8 @@ def read_posed_body(folder, id, poses=None, deorient=False):
     return verts
 
 
-def read_curvatures(h5file, id, frames, neighbourhood=None):
+def read_curvatures(h5file, id, garment, frames, neighbourhood=None):
-    frames_filtered = [frame for frame in frames if frame.split()[0] == id]
+    frames_filtered = [frame for frame in frames if (frame.split()[0] == id and frame.split()[2] == garment)]
     garment = frames_filtered[0].split()[2]
     curvatures_all = []
     for frame_nr in range(len(frames_filtered)):

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)
 
@@ -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'), 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'), exist_ok=True)
+    os.makedirs(os.path.join(extras_folder, 'uv_mask'), exist_ok=True)
 
 
     for id in tqdm.tqdm(ids):
@@ -416,14 +493,15 @@ if __name__ == '__main__':
         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_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(h5file, id, frames, neighbourhood=None)
+            mesh_curvature = read_curvatures(h5file_curv, id, garment, frames, neighbourhood=None)
         if 'body_posed' in data_to_convert:
             body_posed = read_posed_body(folder, id, info['poses'], use_no_orient)
 
@@ -454,35 +532,45 @@ if __name__ == '__main__':
 
         start = time.perf_counter()
 
-        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_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 = 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)
+            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:
-            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)
+            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, garment_uv_mask_border=garment_uv_mask_border).astype(np.float32)
+            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(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)
+            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)
+
+
 
         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_static, 'uv_static', 'float32')
-        save_to_h5file(h5file, id, 0, garment, uv_mask, 'uv_mask', '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)
+
         if 'mesh_posed' in data_to_convert:
-            for i in range(uv_posed.shape[0]):
+            for i in range(len(uv_posed)):
-                save_to_h5file(h5file, id, i, garment, uv_posed[i, :, :, :], 'uv_posed' + orient_mod + shape_mod, 'float32')
+                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(uv_curvature.shape[0]):
+            for i in range(len(uv_curvature)):
-                save_to_h5file(h5file, id, i, garment, uv_curvature[i, :, :, :], 'uv_curvature', 'uint8')
+                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(uv_body_posed.shape[0]):
+            for i in range(len(uv_body_posed)):
-                save_to_h5file(h5file, id, i, garment, uv_body_posed[i, :, :, :], 'uv_body_posed', 'float32')
+                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')
         
 
         # start = time.perf_counter()
@@ -500,9 +588,9 @@ if __name__ == '__main__':
         # 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)
+        # os.makedirs(os.path.join(folder, f'uv_curvature_{neighbourhood}_2'), exist_ok=True)
-        for i in range(uv_curvature.shape[0]):
+        # for i in range(uv_curvature.shape[0]):
-            np.savez_compressed(os.path.join(folder, f'uv_curvature_{neighbourhood}_2', id + '_' + str(i) + '.npz'), curvature_uv=uv_curvature[i, :, :])
+        #     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]):
@@ -529,3 +617,5 @@ if __name__ == '__main__':
 
         pass
 
+
+    h5file_curv.close()

data/curvatures.py

@@ -127,10 +127,10 @@ class Cache:
             if self.include_neighbours:
                 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)
+                # G = tm.graph.vertex_adjacency_graph(m)
-                for i in range(v.shape[0]):
+                # for i in range(v.shape[0]):
-                    for k in range(25):
+                #     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]}'
+                #         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
@@ -331,11 +331,12 @@ 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("-e", "--end", help = "end of range", default=600000)
     parser.add_argument("-n", "--name", help = "end of range", default="lhunor")
     args = parser.parse_args()
 

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:

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 = 0
+    step = 100
+
+    # !!!!!!! change skirt parameter in script
+
+    while count < 2000:
+        # 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_' + str(count) + '_' + str(count+step) + ' > nohup/to_uv/train_' + 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_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 utils.losses import collision_ratio, kl_loss, normal_loss, collision_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 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')
+    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.'))]
-    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.')]
     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:
+        if train_accumulator.get()['train/euclid_error'] < (config['balancing_activation_limit'] if 'balancing_activation_limit' in config.keys() else 0.017):
-            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 balancing_activated_step == 0:
                 balancing_activated_step = step
             g_balancing_coeff = g_balancing_coeff + 1 / 10000     
 
+
+        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()
+
+