updated pipeline, added uv_mask_loss

This commit is contained in:
2022-06-12 15:50:30 +00:00
parent 94e693d7f6
commit def8a85cf9
6 changed files with 160 additions and 78 deletions
+4 -3
View File
@@ -23,15 +23,16 @@ class DatasetSimple(tf.data.Dataset):
uv_static = np.load(os.path.join(folder, 'uv_static', id + '.npz')); uv_static = uv_static['uv_static']; uv_static = np.transpose(uv_static, [2, 0, 1])
mesh_static = np.load(os.path.join(folder, 'mesh_static', id + '.npz')); mesh_static = mesh_static['mesh_static']
mesh_body_posed = np.load(os.path.join(folder, 'mesh_body_posed', id + '_' + frame_nr + '.npy'))
mask = np.load(os.path.join(folder, 'mask', id + '.npy')); mask = mask.astype(np.float32)
mesh_mask = np.load(os.path.join(folder, 'mesh_mask', id + '.npy')); mesh_mask = mesh_mask.astype(np.float32)
uv_mask = np.load(os.path.join(folder, 'uv_mask', id + '.npz')); uv_mask=uv_mask['uv_mask']; uv_mask = uv_mask.astype(np.float32)
# uv_posed, (pose, shape, uv_static), mesh_static, mesh_body_posed, mask
yield uv_posed, (pose_shape[:72], pose_shape[72:], uv_static), mesh_static, mesh_body_posed, mask
yield uv_posed, (pose_shape[:72], pose_shape[72:], uv_static), mesh_static, mesh_body_posed, mesh_mask, uv_mask
def __new__(cls, folder, res=256):
return tf.data.Dataset.from_generator(
cls._generator,
output_signature = (tf.TensorSpec(shape = (3, res * 2, res,), dtype = tf.float32), (tf.TensorSpec(shape = (72), dtype = tf.float32), tf.TensorSpec(shape = (11), dtype = tf.float32), tf.TensorSpec(shape = (3, res * 2, res,), dtype = tf.float32) ), tf.TensorSpec(shape = (14475, 3,), dtype = tf.float32), tf.TensorSpec(shape = (14475, 3,), dtype = tf.float32), tf.TensorSpec(shape = (14475,), dtype = tf.float32) ),
output_signature = (tf.TensorSpec(shape = (3, res * 2, res,), dtype = tf.float32), (tf.TensorSpec(shape = (72), dtype = tf.float32), tf.TensorSpec(shape = (11), dtype = tf.float32), tf.TensorSpec(shape = (3, res * 2, res,), dtype = tf.float32) ), tf.TensorSpec(shape = (14475, 3,), dtype = tf.float32), tf.TensorSpec(shape = (14475, 3,), dtype = tf.float32), tf.TensorSpec(shape = (14475,), dtype = tf.float32), tf.TensorSpec(shape = (res * 2, res,), dtype = tf.float32) ),
args=(folder, )
)