luni/cvpr
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

add use_no_orient switch

Browse Source
This commit is contained in:
luni
2022-10-29 23:46:58 +00:00
parent d59639b8d9
commit 745193dc94
3 changed files with 86 additions and 54 deletions
Download Patch File Download Diff File Expand all files Collapse all files
dataset.py
+38 -31
View File
@@ -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 = []