add use_no_orient switch
This commit is contained in:
+25
-19
@@ -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
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
+34
-27
@@ -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'))
|
||||
|
||||
##############################
|
||||
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():
|
||||
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]
|
||||
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 = []
|
||||
|
||||
Reference in New Issue
Block a user