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

curv by default is one_hot

Browse Source
This commit is contained in:
luni
2022-09-27 21:24:41 +00:00
parent a5eff778df
commit 64264364be
3 changed files with 30 additions and 27 deletions
Download Patch File Download Diff File Expand all files Collapse all files
dataset.py
+18 -12
View File
@@ -61,7 +61,7 @@ class StaticDataset():
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_mask_32', 'uv_mask_128', 'uv_shaped_body']
self.keys = [value for value in signature.keys() if value in self.possible_keys]
if 'uv_static_128' in signature.keys():
if 'uv_static_128' in signature.keys() and 'uv_static_128' not in self.keys:
self.keys.append('uv_static_128')
self.resolutions = resolutions
self.filled = ''
@@ -132,14 +132,14 @@ class DatasetFactory():
for key in signature.keys():
if key == 'id':
tf_signature.append(tf.TensorSpec(shape = (), dtype = tf.string))
elif key == 'uv_curvature' or key == 'mesh_curvature' or key == 'uv_curvature_32' or key == 'uv_curvature_128':
elif key == 'mesh_curvature':
tf_signature.append(tf.TensorSpec(shape = signature[key], dtype = tf.int32))
else:
tf_signature.append(tf.TensorSpec(shape = signature[key], dtype = tf.float32))
return tf_signature
def get_dataset(self, folder, batch_size, signature, shuffle=True, use_filled=False, resolutions=[512]):
def get_dataset(self, folder, batch_size, signature, shuffle=True, use_filled=False, resolutions=[512], in_memory=True):
# assert (use_filled and len(resolutions) >= 1) or (not use_filled and len(resolutions) == 1)
assert set(resolutions).issubset(set([32, 128, 512]))
@@ -212,7 +212,8 @@ class DatasetFactory():
uv_curvature = uv_curvature['uv_curvature' + filled]
else:
uv_curvature = uv_curvature['curvature_uv']
item['uv_curvature' + res_string] = uv_curvature.astype(np.int32)
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' + filled + res_string, id + '_' + frame_nr + '.npz')); uv_unposed_no_shape = uv_unposed_no_shape['uv_unposed_no_shape' + filled]
@@ -235,10 +236,18 @@ class DatasetFactory():
# inp=[i],
# Tout=self._build_tf_signature(signature)
# ), num_parallel_calls=tf.data.AUTOTUNE).cache().batch(batch_size=batch_size, drop_remainder=True).repeat()
dataset = dataset_idx.map(lambda i: tf.py_function(func=read_item,
inp=[i],
Tout=self._build_tf_signature(signature)
), num_parallel_calls=tf.data.AUTOTUNE).cache().batch(batch_size=batch_size, drop_remainder=True).repeat()
if in_memory:
dataset = dataset_idx.map(lambda i: tf.py_function(func=read_item,
inp=[i],
Tout=self._build_tf_signature(signature)
), num_parallel_calls=tf.data.AUTOTUNE).cache().batch(batch_size=batch_size, drop_remainder=True).repeat()
else:
dataset = dataset_idx.map(lambda i: tf.py_function(func=read_item,
inp=[i],
Tout=self._build_tf_signature(signature)
), num_parallel_calls=tf.data.AUTOTUNE).batch(batch_size=batch_size, drop_remainder=True).repeat()
# dataset = dataset_idx.map(read_item, num_parallel_calls=tf.data.AUTOTUNE).cache().batch(batch_size=batch_size, drop_remainder=True).repeat()
return dataset
@@ -281,7 +290,7 @@ if __name__ == '__main__':
# 'mesh_unposed': (14475, 3,),
#'uv_unposed': (3, 512, 256,),
'uv_unposed_32': (3, 32, 16,),
# 'uv_curvature': (512, 256),
'uv_curvature': (512, 256, 7),
# 'mesh_curvature': (14475)
}
signature_static = { #'uv_static': (3, 512, 256,),
@@ -335,9 +344,6 @@ if __name__ == '__main__':
print()
uv_res = 256
def m2uv_array(m2uv):