diff --git a/configs/config_cvae_test.py b/configs/config_cvae_test.py index a3c58be..e063baf 100644 --- a/configs/config_cvae_test.py +++ b/configs/config_cvae_test.py @@ -7,31 +7,43 @@ config = dict() # DON'T FORGET TO LOAD FINAL UPSCALER UNLESS HAVE LEVEL 512 # DON'T FORGET TO CHECK AVG GARMENT USED +config['use_skinning'] = True # and deshaping +config['remove_global_orientation'] = False +config['use_curvature'] = True +config['use_extra_losses'] = True # edge, collision, 3d mesh +config['use_offsets'] = True + +assert [config['use_skinning'], config['remove_global_orientation'], config['use_curvature'], config['use_extra_losses'], config['use_offsets']].count(False) == 1 + +config['gpu'] = '2' +config['name'] = 'ab_1' +config['notes'] = '' + +config['avg_garment_path'] = 'custom_data/train_no_shape_avg_uv.npy' # train_no_shape_avg_uv train_no_shape_one_hot_avg_uv +config['reg_loss_coeff'] = 0.1 # set to 0 if don't want to use edge+collision + config['batch_size'] = 4 -config['train_dir'] = '/home/workspace/local-data/train' -config['valid_dir'] = '/home/workspace/local-data/valid' -config['test_dir'] = '/home/workspace/data/dataset/test_ready' -config['use_filled'] = False config['resolution'] = 512 config['is_skirt'] = False config['dataset_in_memory'] = False config['is_cape'] = False +config['train_dir'] = '/home/workspace/local-data/train' +config['valid_dir'] = '/home/workspace/local-data/valid' +config['test_dir'] = '/home/workspace/data/dataset/test_ready' + +config['static_model'] = 'checkpoints/static_1' +config['body_model'] = 'checkpoints/body_1' +config['curvature_model'] = 'checkpoints/curv_1' + + config['smpl_res'] = 14475 if config['is_cape']: config['smpl_res'] = 3638 config['latent_dim'] = 128 - -config['gpu'] = '1' -config['name'] = 'ab_1' -config['notes'] = '' -config['pretrained_model'] = 'checkpoints/v2_5_2' - -config['resolution'] = 512 -config['is_skirt'] = False -config['reg_loss_coeff'] = 0.1 # set to 0 if don't want to use edge+collision +config['pretrained_model'] = '' config['l2_coeff'] = 1 ############### @@ -53,12 +65,6 @@ config['signature_static'] = { 'uv_static': (512, 256, 3), 'uv_mask': (512, 256,), } -config['avg_garment_path'] = 'custom_data/train_no_shape_avg_uv.npy' # train_no_shape_avg_uv train_no_shape_one_hot_avg_uv - - -config['static_model'] = 'checkpoints/static_1' -config['body_model'] = 'checkpoints/body_1' -config['curvature_model'] = 'checkpoints/curv_1' config['upscale_enhancer'] = '' config['freeze_encoders'] = True config['use_one_hot'] = False diff --git a/cvae.py b/cvae.py index 909ad90..4776857 100644 --- a/cvae.py +++ b/cvae.py @@ -427,10 +427,29 @@ if __name__ == '__main__': factory = DatasetFactory() - dataset_train = factory.get_dataset(folder=config['train_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'], use_filled=config['use_filled'], resolutions=[config['resolution']], in_memory=config['dataset_in_memory'], use_one_hot=config['use_one_hot']) - static_dataset_train = StaticDataset(folder=config['train_dir'], signature=config['signature_static'], use_filled=config['use_filled'], resolutions=[config['resolution']], in_memory=True) - dataset_valid = factory.get_dataset(folder=config['valid_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'], use_filled=config['use_filled'], resolutions=[config['resolution']], in_memory=config['dataset_in_memory'], use_one_hot=config['use_one_hot']) - static_dataset_valid = StaticDataset(folder=config['valid_dir'], signature=config['signature_static'], use_filled=config['use_filled'], resolutions=[config['resolution']], in_memory=True) + dataset_train = factory.get_dataset(folder=config['train_dir'], + batch_size=config['batch_size'], + signature=config['signature_dynamic'], + resolutions=[config['resolution']], + in_memory=config['dataset_in_memory'], + use_one_hot=config['use_one_hot'], + use_no_orient=config['remove_global_orientation']) + static_dataset_train = StaticDataset(folder=config['train_dir'], + signature=config['signature_static'], + resolutions=[config['resolution']], + in_memory=True) + + dataset_valid = factory.get_dataset(folder=config['valid_dir'], + batch_size=config['batch_size'], + signature=config['signature_dynamic'], + resolutions=[config['resolution']], + in_memory=config['dataset_in_memory'], + use_one_hot=config['use_one_hot'], + use_no_orient=config['remove_global_orientation']) + static_dataset_valid = StaticDataset(folder=config['valid_dir'], + signature=config['signature_static'], + resolutions=[config['resolution']], + in_memory=True) model = CVAE(config) diff --git a/dataset.py b/dataset.py index 1478cb5..ad7a937 100644 --- a/dataset.py +++ b/dataset.py @@ -167,15 +167,21 @@ class DatasetFactory(): return tf_signature - def get_dataset(self, folder, batch_size, signature, shuffle=True, use_filled=False, resolutions=[512], in_memory=False, do_repeat=True, is_cape=False, use_one_hot=False): - # assert (use_filled and len(resolutions) >= 1) or (not use_filled and len(resolutions) == 1) + def get_dataset(self, + folder, + batch_size, + signature, + shuffle=True, + resolutions=[512], + in_memory=False, + do_repeat=True, + is_cape=False, + use_one_hot=False, + use_no_orient=True): assert set(resolutions).issubset(set([32, 64, 128, 256, 512])) - # static_data = StaticDataset(folder, signature) dataset_idx = DatasetIdx(folder, shuffle, is_cape=is_cape) - filled = '' - if use_filled: - filled = '_filled' + separator = '_' if is_cape: separator = '.' @@ -186,6 +192,13 @@ class DatasetFactory(): one_hot = '_one_hot' one_hot_2 = '_one_hot_2' + uv_posed_body_no_orient_folder_name = 'uv_posed_body_no_orient_2' + if is_cape: + uv_posed_body_no_orient_folder_name = 'uv_posed_body_no_orient' + + if use_no_orient == False: + assert resolutions == [512], 'Using global orientation is only possible with non-pyramid model!' + def read_item(frame): frame_split = tf.strings.split(frame, sep=separator) @@ -198,28 +211,19 @@ class DatasetFactory(): item['id'] = id if 'pose_shape' in signature.keys(): pose_shape = np.load(os.path.join(folder, 'pose_shape', id + separator + frame_nr + '.npy')) - - ############################## - pose_shape[:3] = 0 - ############################## - + if use_no_orient: + pose_shape[:3] = 0 item['pose_shape'] = pose_shape if 'uv_posed' in signature.keys(): uv_posed = np.load(os.path.join(folder, 'uv_posed', id + separator + frame_nr + '.npz')); uv_posed = uv_posed['uv_gt'] item['uv_posed'] = uv_posed - # if 'uv_body_posed' in signature.keys(): - # uv_body_posed = np.load(os.path.join(folder, 'uv_body_posed', id + '_' + frame_nr + '.npz')); uv_body_posed = uv_body_posed['uv_body']; uv_body_posed = np.transpose(uv_body_posed, [2, 0, 1]) - # item['uv_body_posed'] = uv_body_posed - # if 'mesh_body_posed' in signature.keys(): - # mesh_body_posed = np.load(os.path.join(folder, 'mesh_body_posed', id + '_' + frame_nr + '.npy')) - # item['mesh_body_posed'] = mesh_body_posed - # if 'mesh_posed_body_no_orient' in signature.keys(): - # mesh_posed_body_no_orient = np.load(os.path.join(folder, 'mesh_posed_body_no_orient', id + '_' + frame_nr + '.npz')); mesh_posed_body_no_orient = mesh_posed_body_no_orient['mesh_posed_body_no_orient'] - # item['mesh_posed_body_no_orient'] = mesh_posed_body_no_orient - if 'mesh_body_posed' in signature.keys(): - mesh_posed_body_no_orient = np.load(os.path.join(folder, 'mesh_posed_body_no_orient', id + separator + frame_nr + '.npz')); mesh_posed_body_no_orient = mesh_posed_body_no_orient['mesh_posed_body_no_orient'] - item['mesh_body_posed'] = mesh_posed_body_no_orient + if use_no_orient: + mesh_posed_body_no_orient = np.load(os.path.join(folder, 'mesh_posed_body_no_orient', id + separator + frame_nr + '.npz')); mesh_posed_body_no_orient = mesh_posed_body_no_orient['mesh_posed_body_no_orient'] + item['mesh_body_posed'] = mesh_posed_body_no_orient + else: + mesh_body_posed = np.load(os.path.join(folder, 'mesh_body_posed', id + '_' + frame_nr + '.npy')) + item['mesh_body_posed'] = mesh_body_posed if 'mesh_posed' in signature.keys(): mesh_posed = np.load(os.path.join(folder, 'mesh_posed', id + separator + frame_nr + '.npz')); mesh_posed = mesh_posed['mesh_posed'] item['mesh_posed'] = mesh_posed @@ -246,23 +250,26 @@ class DatasetFactory(): res_string = '_' + str(res) if 'uv_curvature' + res_string in signature.keys(): - uv_curvature = np.load(os.path.join(folder, 'uv_curvature' + filled + res_string, id + separator + frame_nr + '.npz')); - if filled != '' or res_string != '': - uv_curvature = uv_curvature['uv_curvature' + filled] + uv_curvature = np.load(os.path.join(folder, 'uv_curvature' + res_string, id + separator + frame_nr + '.npz')); + if res_string != '': + uv_curvature = uv_curvature['uv_curvature'] else: uv_curvature = uv_curvature['curvature_uv'] uv_curvature = tf.one_hot(uv_curvature, 7) item['uv_curvature' + res_string] = uv_curvature if 'uv_unposed' + res_string in signature.keys(): - uv_unposed_no_shape = np.load(os.path.join(folder, 'uv_unposed_no_shape' + one_hot_2 + res_string, id + separator + frame_nr + '.npz')); uv_unposed_no_shape = uv_unposed_no_shape['uv_unposed_no_shape' + filled] + uv_unposed_no_shape = np.load(os.path.join(folder, 'uv_unposed_no_shape' + one_hot_2 + res_string, id + separator + frame_nr + '.npz')); uv_unposed_no_shape = uv_unposed_no_shape['uv_unposed_no_shape'] item['uv_unposed' + res_string] = uv_unposed_no_shape if 'uv_body_posed' + res_string in signature.keys(): - uv_posed_body_no_orient = np.load(os.path.join(folder, 'uv_posed_body_no_orient_2' + filled + res_string, id + separator + frame_nr + '.npz')) - uv_posed_body_no_orient = uv_posed_body_no_orient['uv_posed_body_no_orient' + filled] - item['uv_body_posed' + res_string] = uv_posed_body_no_orient - + if use_no_orient: + uv_posed_body_no_orient = np.load(os.path.join(folder, uv_posed_body_no_orient_folder_name + res_string, id + separator + frame_nr + '.npz')) + uv_posed_body_no_orient = uv_posed_body_no_orient['uv_posed_body_no_orient'] + item['uv_body_posed' + res_string] = uv_posed_body_no_orient + else: + uv_body_posed = np.load(os.path.join(folder, 'uv_body_posed', id + '_' + frame_nr + '.npz')); uv_body_posed = uv_body_posed['uv_body'] + item['uv_body_posed'] = uv_body_posed # item = { **static_info, **item } as_list = []