added one_hot option to dataset

This commit is contained in:
2022-10-29 20:09:46 +00:00
parent db3ea8ad48
commit 2e8d6f04f2
+21 -10
View File
@@ -167,7 +167,7 @@ class DatasetFactory():
return tf_signature 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): 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) # assert (use_filled and len(resolutions) >= 1) or (not use_filled and len(resolutions) == 1)
assert set(resolutions).issubset(set([32, 64, 128, 256, 512])) assert set(resolutions).issubset(set([32, 64, 128, 256, 512]))
@@ -180,6 +180,12 @@ class DatasetFactory():
if is_cape: if is_cape:
separator = '.' separator = '.'
one_hot = ''
one_hot_2 = ''
if use_one_hot:
one_hot = '_one_hot'
one_hot_2 = '_one_hot_2'
def read_item(frame): def read_item(frame):
frame_split = tf.strings.split(frame, sep=separator) frame_split = tf.strings.split(frame, sep=separator)
@@ -224,7 +230,8 @@ class DatasetFactory():
# uv_unposed = np.load(os.path.join(folder, 'uv_unposed', id + '_' + frame_nr + '.npz')); uv_unposed = uv_unposed['uv_unposed']; uv_unposed = np.transpose(uv_unposed, [2, 0, 1]) # uv_unposed = np.load(os.path.join(folder, 'uv_unposed', id + '_' + frame_nr + '.npz')); uv_unposed = uv_unposed['uv_unposed']; uv_unposed = np.transpose(uv_unposed, [2, 0, 1])
# item['uv_unposed'] = uv_unposed # item['uv_unposed'] = uv_unposed
if 'mesh_unposed' in signature.keys(): if 'mesh_unposed' in signature.keys():
mesh_unposed_no_shape = np.load(os.path.join(folder, 'mesh_unposed_no_shape', id + separator + frame_nr + '.npz')); mesh_unposed_no_shape = mesh_unposed_no_shape['mesh_unposed_no_shape'] # print('one_hot')
mesh_unposed_no_shape = np.load(os.path.join(folder, 'mesh_unposed_no_shape' + one_hot, id + separator + frame_nr + '.npz')); mesh_unposed_no_shape = mesh_unposed_no_shape['mesh_unposed_no_shape']
item['mesh_unposed'] = mesh_unposed_no_shape item['mesh_unposed'] = mesh_unposed_no_shape
@@ -248,11 +255,12 @@ class DatasetFactory():
item['uv_curvature' + res_string] = uv_curvature item['uv_curvature' + res_string] = uv_curvature
if 'uv_unposed' + res_string in signature.keys(): if 'uv_unposed' + res_string in signature.keys():
uv_unposed_no_shape = np.load(os.path.join(folder, 'uv_unposed_no_shape' + filled + 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' + filled]
item['uv_unposed' + res_string] = uv_unposed_no_shape item['uv_unposed' + res_string] = uv_unposed_no_shape
if 'uv_body_posed' + res_string in signature.keys(): 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' + filled + res_string, id + separator + frame_nr + '.npz')); uv_posed_body_no_orient = uv_posed_body_no_orient['uv_posed_body_no_orient' + filled] 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 item['uv_body_posed' + res_string] = uv_posed_body_no_orient
@@ -318,8 +326,8 @@ if __name__ == '__main__':
signature_dynamic = { 'id': (), signature_dynamic = { 'id': (),
# 'pose_shape': (72 + 11,), # 'pose_shape': (72 + 11,),
# 'uv_posed': (3, 512, 256,), # 'uv_posed': (3, 512, 256,),
#'uv_body_posed': (3, 512, 256,), # 'uv_body_posed': (512, 256, 3),
# 'uv_body_posed_32': (3, 32, 16,), # 'uv_body_posed_32': (32, 16, 3),
# 'mesh_body_posed': (14475, 3,), # 'mesh_body_posed': (14475, 3,),
# 'mesh_posed': (14475, 3,), # 'mesh_posed': (14475, 3,),
# 'mesh_unposed': (14475, 3,), # 'mesh_unposed': (14475, 3,),
@@ -341,13 +349,16 @@ if __name__ == '__main__':
signature=signature_dynamic, signature=signature_dynamic,
use_filled=False, use_filled=False,
resolutions=[512], resolutions=[512],
do_repeat=False) do_repeat=False,
use_one_hot=True)
static_dataset = StaticDataset(folder='/home/workspace/local-data/train', static_dataset = StaticDataset(folder='/home/workspace/local-data/train',
signature=signature_static, signature=signature_static,
use_filled=False, use_filled=False,
resolutions=[512]) resolutions=[512])
# for i, batch in enumerate(dataset):
# print(i)
sum_uv = np.zeros((512, 256, 3), dtype=np.float32) sum_uv = np.zeros((512, 256, 3), dtype=np.float32)
@@ -375,11 +386,11 @@ if __name__ == '__main__':
k = F[:,2] k = F[:,2]
)]).show() )]).show()
np.save('custom_data/train_no_shape_avg_uv.npy', avg_garment) np.save('custom_data/train_no_shape_one_hot_avg_uv.npy', avg_garment)
print() # print()
pass # pass
# uv_res = 256 # uv_res = 256