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

train encoder on raw curvature values

Browse Source
This commit is contained in:
luni
2023-01-27 11:13:10 +00:00
parent a3c5b526b8
commit 6e6b0a6604
8 changed files with 60 additions and 71 deletions
Download Patch File Download Diff File Expand all files Collapse all files
dataset.py
+13 -14
View File
@@ -176,7 +176,7 @@ class DatasetFactory():
elif key == 'frame_nr':
tf_signature.append(tf.TensorSpec(shape = (), dtype = tf.string))
elif key == 'mesh_curvature':
tf_signature.append(tf.TensorSpec(shape = signature[key], dtype = tf.int32))
tf_signature.append(tf.TensorSpec(shape = signature[key], dtype = tf.float32))
else:
tf_signature.append(tf.TensorSpec(shape = signature[key], dtype = tf.float32))
return tf_signature
@@ -260,8 +260,8 @@ class DatasetFactory():
if 'mesh_curvature' in signature.keys():
mesh_curvature = np.load(os.path.join(folder, 'mesh_curvature', id + separator + frame_nr + '.npz')); mesh_curvature = mesh_curvature['mesh_curvature']
item['mesh_curvature'] = mesh_curvature.astype(np.int32)
mesh_curvature = np.load(os.path.join(folder, 'mesh_curvature_raw', id + separator + frame_nr + '.npz')); mesh_curvature = mesh_curvature['mesh_curvature_raw']
item['mesh_curvature'] = mesh_curvature
for res in resolutions:
@@ -270,13 +270,13 @@ class DatasetFactory():
res_string = '_' + str(res)
if 'uv_curvature' + res_string in signature.keys():
uv_curvature = np.load(os.path.join(folder, 'uv_curvature' + res_string, id + separator + frame_nr + '.npz'));
uv_curvature = np.load(os.path.join(folder, 'curvature_uv_raw' + 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
uv_curvature = uv_curvature['curvature_uv_raw']
# uv_curvature = tf.one_hot(uv_curvature, 7)
item['uv_curvature' + res_string] = uv_curvature[:,:,tf.newaxis]
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']
@@ -370,8 +370,8 @@ if __name__ == '__main__':
# 'mesh_unposed': (14475, 3,),
'uv_unposed': (512, 256, 3),
# 'uv_unposed_32': (3, 32, 16,),
'uv_curvature': (512, 256, 7),
# 'mesh_curvature': (14475)
'uv_curvature': (512, 256, 1),
'mesh_curvature': (14475)
}
signature_static = { #'uv_static': (3, 512, 256,),
# 'uv_static_32': (3, 32, 16,),
@@ -381,15 +381,14 @@ if __name__ == '__main__':
# 'uv_shaped_body': (3, 512, 256,),
}
dataset = factory.get_dataset(folder='/home/workspace/local-data/train_skirt',
dataset = factory.get_dataset(folder='/home/workspace/local-data/train',
batch_size=1,
signature=signature_dynamic,
resolutions=[512],
do_repeat=False,
use_one_hot=False,
is_cape=False,
frames_file='frames_dress_30k.txt')
static_dataset = StaticDataset(folder='/home/workspace/local-data/train_skirt',
is_cape=False,)
static_dataset = StaticDataset(folder='/home/workspace/local-data/train',
signature=signature_static,
resolutions=[512],
is_cape=False)
@@ -404,7 +403,7 @@ if __name__ == '__main__':
for i, batch in enumerate(dataset):
print(i)
batch = static_dataset.expand_batch(batch_to_dict(batch, signature_dynamic))
ex.imshow(batch['uv_curvature'][0].numpy().argmax(-1) / 6.0).show()
ex.imshow(batch['uv_curvature'][0,:,:,0].numpy()).show()
sum_uv += batch['uv_unposed'].numpy().sum(axis=0)
sum_mask += batch['uv_mask'].numpy().sum(axis=0)