From a94be773d3380f2f378cdd9c23ae032d4208795f Mon Sep 17 00:00:00 2001 From: Hunor Laczko Date: Wed, 15 Feb 2023 16:21:14 +0000 Subject: [PATCH] save progress --- .gitignore | 1 + compression_test.py | 32 ++++ custom_data/train_bins_25.npy | Bin 184 -> 184 bytes data/cloth3d_io.py | 11 +- data/convert_to_uv.py | 132 ++++++++++++++--- data/curvatures.py | 11 +- data/preprocessing.py | 4 +- run_parallel.py | 29 ++++ scripts/cvae.py | 76 +++++----- utils/losses.py | 270 ++++++++++++++++++++++++++++++++++ utils/postprocess.py | 137 +++++++++++++++++ utils/utils.py | 27 ++++ 12 files changed, 656 insertions(+), 74 deletions(-) create mode 100644 compression_test.py create mode 100644 run_parallel.py create mode 100644 utils/losses.py create mode 100644 utils/postprocess.py create mode 100644 utils/utils.py diff --git a/.gitignore b/.gitignore index 745e0ec..ad94b87 100644 --- a/.gitignore +++ b/.gitignore @@ -172,3 +172,4 @@ ids/ examples/ custom_data/dijkstra* .vscode/* +.vscode/launch.json diff --git a/compression_test.py b/compression_test.py new file mode 100644 index 0000000..4857c61 --- /dev/null +++ b/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')) \ No newline at end of file diff --git a/custom_data/train_bins_25.npy b/custom_data/train_bins_25.npy index c5c2a3ca4b2fe39509d8521e0830ef7e8894fdde..9ed5faf824555537455fa569b3c4ab4fc97f7c80 100644 GIT binary patch delta 60 ncmdnNxPx&*v#EpBie>v4aD#>wul(#m(#huaAo{`OhAZ{}VHONt delta 60 zcmdnNxPx&*v#Emg?iu?S5MY7n{3v@M*TGh+97HFaztRJu89Z(+1dCq_xMB|g&qWWq diff --git a/data/cloth3d_io.py b/data/cloth3d_io.py index 572210b..52afa52 100644 --- a/data/cloth3d_io.py +++ b/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')) - 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, :] @@ -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): - frames_filtered = [frame for frame in frames if frame.split()[0] == id] +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)] garment = frames_filtered[0].split()[2] curvatures_all = [] for frame_nr in range(len(frames_filtered)): diff --git a/data/convert_to_uv.py b/data/convert_to_uv.py index f9e5baf..7790a32 100644 --- a/data/convert_to_uv.py +++ b/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_mask, 'uv_mask', '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') + 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]): - save_to_h5file(h5file, id, i, garment, uv_posed[i, :, :, :], 'uv_posed' + orient_mod + shape_mod, 'float32') + 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(uv_curvature.shape[0]): - save_to_h5file(h5file, id, i, garment, uv_curvature[i, :, :, :], 'uv_curvature', 'uint8') + 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(uv_body_posed.shape[0]): - save_to_h5file(h5file, id, i, garment, uv_body_posed[i, :, :, :], 'uv_body_posed', 'float32') + 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') # 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) - 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, 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]): @@ -529,3 +617,5 @@ if __name__ == '__main__': pass + + h5file_curv.close() \ No newline at end of file diff --git a/data/curvatures.py b/data/curvatures.py index 45d59ab..846e04b 100644 --- a/data/curvatures.py +++ b/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) - 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]}' + # 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 @@ -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() diff --git a/data/preprocessing.py b/data/preprocessing.py index 7b6069d..53cd1c5 100644 --- a/data/preprocessing.py +++ b/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: diff --git a/run_parallel.py b/run_parallel.py new file mode 100644 index 0000000..d45bc3d --- /dev/null +++ b/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 diff --git a/scripts/cvae.py b/scripts/cvae.py index b66874a..27637a3 100644 --- a/scripts/cvae.py +++ b/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 + g_balancing_coeff = g_balancing_coeff + 1 / 10000 - to_log = { **to_log, **{ 'train/balancing_coeff': g_balancing_coeff } } + + 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) + wandb.log(to_log, step = step) + train_accumulator.reset() + valid_accumulator.reset() + + if step % 5000 == 0: + manager.save(checkpoint_number=step) print(step) diff --git a/utils/losses.py b/utils/losses.py new file mode 100644 index 0000000..430359e --- /dev/null +++ b/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() diff --git a/utils/postprocess.py b/utils/postprocess.py new file mode 100644 index 0000000..e1d3da3 --- /dev/null +++ b/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() \ No newline at end of file diff --git a/utils/utils.py b/utils/utils.py new file mode 100644 index 0000000..337bc40 --- /dev/null +++ b/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() + +