train on template only with mask applied

2022-06-14 10:16:55 +00:00
parent def8a85cf9
commit 245f811d47
11 changed files with 112 additions and 60 deletions
@@ -27,12 +27,12 @@ class DatasetSimple(tf.data.Dataset):
            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, mesh_mask, uv_mask
+            yield uv_posed, (pose_shape[:72], pose_shape[72:], uv_static, uv_mask), mesh_static, mesh_body_posed, mesh_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), tf.TensorSpec(shape = (res * 2, res,), 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 = (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) ),
            args=(folder, )
        )