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

add handling of different resolutions

Browse Source
This commit is contained in:
luni
2022-09-13 15:45:51 +00:00
parent 701bd525d3
commit 6ca873c76e
1 changed files with 56 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files
dataset.py
+56 -26
View File
@@ -54,22 +54,31 @@ class DatasetSimple(tf.data.Dataset):
class StaticDataset(): class StaticDataset():
def __init__(self, folder, signature, use_filled=True) -> None: def __init__(self, folder, signature, use_filled=True, resolutions=[512]) -> None:
self.store = dict() self.store = dict()
self.signature = signature self.signature = signature
self.folder = folder self.folder = folder
self.possible_keys = ['uv_static', 'mesh_static', 'mesh_mask', 'uv_mask', 'uv_shaped_body'] assert set(resolutions).issubset(set([32, 128, 512]))
self.possible_keys = ['uv_static', 'uv_static_32', 'uv_static_128' 'mesh_static', 'mesh_mask', 'uv_mask', 'uv_shaped_body']
self.keys = [value for value in signature.keys() if value in self.possible_keys] self.keys = [value for value in signature.keys() if value in self.possible_keys]
filled = '' self.resolutions = resolutions
self.filled = ''
if use_filled: if use_filled:
filled = '_filled' self.filled = '_filled'
def read(self, id): def read(self, id):
data = dict() data = dict()
if 'uv_static' in self.signature.keys(): for res in self.resolutions:
uv_static = np.load(os.path.join(self.folder, 'uv_static' + filled, id + '.npz')); uv_static = uv_static['uv_static']; uv_static = np.transpose(uv_static, [2, 0, 1]); data['uv_static' + filled]=uv_static res_string = ''
if res != 512:
res_string = '_' + str(res)
if 'uv_static' + res_string in self.signature.keys():
uv_static = np.load(os.path.join(self.folder, 'uv_static' + self.filled + res_string, id + '.npz')); uv_static = uv_static['uv_static' + self.filled]; uv_static = np.transpose(uv_static, [2, 0, 1]); data['uv_static' + res_string] = uv_static
if 'mesh_static' in self.signature.keys(): if 'mesh_static' in self.signature.keys():
mesh_static = np.load(os.path.join(self.folder, 'mesh_static', id + '.npz')); mesh_static = mesh_static['mesh_static']; data['mesh_static'] = mesh_static mesh_static = np.load(os.path.join(self.folder, 'mesh_static', id + '.npz')); mesh_static = mesh_static['mesh_static']; data['mesh_static'] = mesh_static
if 'mesh_mask' in self.signature.keys(): if 'mesh_mask' in self.signature.keys():
@@ -124,14 +133,17 @@ class DatasetFactory():
for key in signature.keys(): for key in signature.keys():
if key == 'id': if key == 'id':
tf_signature.append(tf.TensorSpec(shape = (), dtype = tf.string)) tf_signature.append(tf.TensorSpec(shape = (), dtype = tf.string))
elif key == 'uv_curvature' or key == 'mesh_curvature': elif key == 'uv_curvature' or key == 'mesh_curvature' or key == 'uv_curvature_32' or key == 'uv_curvature_128':
tf_signature.append(tf.TensorSpec(shape = signature[key], dtype = tf.int32)) tf_signature.append(tf.TensorSpec(shape = signature[key], dtype = tf.int32))
else: else:
tf_signature.append(tf.TensorSpec(shape = signature[key], dtype = tf.float32)) tf_signature.append(tf.TensorSpec(shape = signature[key], dtype = tf.float32))
return tf_signature return tf_signature
def get_dataset(self, folder, batch_size, signature, shuffle=True, use_filled=True): def get_dataset(self, folder, batch_size, signature, shuffle=True, use_filled=True, resolutions=[512]):
assert (use_filled and len(resolutions) >= 1) or (not use_filled and len(resolutions) == 1)
assert set(resolutions).issubset(set([32, 128, 512]))
# static_data = StaticDataset(folder, signature) # static_data = StaticDataset(folder, signature)
dataset_idx = DatasetIdx(folder, shuffle) dataset_idx = DatasetIdx(folder, shuffle)
filled = '' filled = ''
@@ -167,9 +179,7 @@ class DatasetFactory():
# if 'mesh_posed_body_no_orient' in signature.keys(): # 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'] # 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 # item['mesh_posed_body_no_orient'] = mesh_posed_body_no_orient
if 'uv_body_posed' in signature.keys():
uv_posed_body_no_orient = np.load(os.path.join(folder, 'uv_posed_body_no_orient' + filled, id + '_' + 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.transpose(uv_posed_body_no_orient, [2, 0, 1])
item['uv_body_posed'] = uv_posed_body_no_orient
if 'mesh_body_posed' in signature.keys(): if 'mesh_body_posed' 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'] 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_body_posed'] = mesh_posed_body_no_orient item['mesh_body_posed'] = mesh_posed_body_no_orient
@@ -185,20 +195,35 @@ class DatasetFactory():
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 + '_' + frame_nr + '.npz')); mesh_unposed_no_shape = mesh_unposed_no_shape['mesh_unposed_no_shape'] mesh_unposed_no_shape = np.load(os.path.join(folder, 'mesh_unposed_no_shape', id + '_' + 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
if 'uv_unposed' in signature.keys():
uv_unposed_no_shape = np.load(os.path.join(folder, 'uv_unposed_no_shape' + filled, id + '_' + frame_nr + '.npz')); uv_unposed_no_shape = uv_unposed_no_shape['uv_unposed_no_shape' + filled]; uv_unposed_no_shape = np.transpose(uv_unposed_no_shape, [2, 0, 1])
item['uv_unposed'] = uv_unposed_no_shape
if 'uv_curvature' in signature.keys():
uv_curvature = np.load(os.path.join(folder, 'uv_curvature' + filled, id + '_' + frame_nr + '.npz'));
if filled != '':
uv_curvature = uv_curvature['uv_curvature' + filled]
else:
uv_curvature = uv_curvature['curvature_uv']
item['uv_curvature'] = uv_curvature.astype(np.int32)
if 'mesh_curvature' in signature.keys(): if 'mesh_curvature' in signature.keys():
mesh_curvature = np.load(os.path.join(folder, 'mesh_curvature', id + '_' + frame_nr + '.npz')); mesh_curvature = mesh_curvature['mesh_curvature'] mesh_curvature = np.load(os.path.join(folder, 'mesh_curvature', id + '_' + frame_nr + '.npz')); mesh_curvature = mesh_curvature['mesh_curvature']
item['mesh_curvature'] = mesh_curvature.astype(np.int32) item['mesh_curvature'] = mesh_curvature.astype(np.int32)
for res in resolutions:
res_string = ''
if res != 512:
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 + '_' + frame_nr + '.npz'));
if filled != '':
uv_curvature = uv_curvature['uv_curvature' + filled]
else:
uv_curvature = uv_curvature['curvature_uv']
item['uv_curvature' + res_string] = uv_curvature.astype(np.int32)
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 + '_' + frame_nr + '.npz')); uv_unposed_no_shape = uv_unposed_no_shape['uv_unposed_no_shape' + filled]; uv_unposed_no_shape = np.transpose(uv_unposed_no_shape, [2, 0, 1])
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' + filled + res_string, id + '_' + 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.transpose(uv_posed_body_no_orient, [2, 0, 1])
item['uv_body_posed' + res_string] = uv_posed_body_no_orient
# item = { **static_info, **item } # item = { **static_info, **item }
as_list = [] as_list = []
for key in signature.keys(): for key in signature.keys():
@@ -250,15 +275,18 @@ 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': (3, 512, 256,),
'uv_body_posed_32': (3, 32, 16,),
# '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,),
'uv_unposed': (3, 512, 256,), #'uv_unposed': (3, 512, 256,),
'uv_unposed_32': (3, 32, 16,),
# 'uv_curvature': (512, 256), # 'uv_curvature': (512, 256),
# 'mesh_curvature': (14475) # 'mesh_curvature': (14475)
} }
signature_static = { #'uv_static': (3, 512, 256,), signature_static = { #'uv_static': (3, 512, 256,),
'uv_static_32': (3, 32, 16,),
# 'mesh_static': (14475, 3,), # 'mesh_static': (14475, 3,),
# 'mesh_mask': (14475,), # 'mesh_mask': (14475,),
'uv_mask': (512, 256,), 'uv_mask': (512, 256,),
@@ -266,12 +294,14 @@ if __name__ == '__main__':
} }
dataset = factory.get_dataset(folder='/home/workspace/local-data/train_100', dataset = factory.get_dataset(folder='/home/workspace/local-data/train_100',
batch_size=300, batch_size=4,
signature=signature_dynamic, signature=signature_dynamic,
use_filled=False) use_filled=True,
resolutions=[32])
static_dataset = StaticDataset(folder='/home/workspace/local-data/train_100', static_dataset = StaticDataset(folder='/home/workspace/local-data/train_100',
signature=signature_static, signature=signature_static,
use_filled=False) use_filled=True,
resolutions=[32])