def test_create_model_input_with_forced_off_masks(self):
# Shape = [2, 13, 2].
keypoints_2d = tf.ones_like([2, 13, 2], dtype=tf.float32)
# Shape = [2, 13].
keypoint_masks_2d = tf.constant(
[[1.0, 0.0, 1.0, 0.0, 1.0, 0.0, 1.0, 0.0, 1.0, 0.0, 1.0, 0.0, 1.0],
[1.0, 1.0, 1.0, 0.0, 0.0, 0.0, 1.0, 1.0, 1.0, 0.0, 0.0, 0.0,
1.0]])
# Shape = [2, 39].
_, side_outputs = input_generator.create_model_input(
keypoints_2d,
keypoint_masks_2d,
keypoints_3d=None,
model_input_keypoint_type=common.
MODEL_INPUT_KEYPOINT_TYPE_2D_INPUT,
model_input_keypoint_mask_type=(
common.
MODEL_INPUT_KEYPOINT_MASK_TYPE_MASK_KEYPOINTS_AND_AS_INPUT),
normalize_keypoints_2d=False,
keypoint_profile_2d=keypoint_profiles.Std13KeypointProfile2D(),
keypoint_dropout_probs=(0.5, 0.5),
forced_mask_off_part_names=['HEAD', 'LEFT_SHOULDER', 'LEFT_ANKLE'],
seed=0)
self.assertCountEqual(
side_outputs.keys(),
{'preprocessed_keypoints_2d', 'preprocessed_keypoint_masks_2d'})
expected_preprocessed_keypoint_masks_2d = np.array(
[[0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0],
[0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0, 1.0, 1.0, 0.0, 0.0, 0.0,
1.0]])
self.assertAllClose(side_outputs['preprocessed_keypoint_masks_2d'],
expected_preprocessed_keypoint_masks_2d)
def testReadFromTable(self):
testdata_dir = 'poem/testdata' # Assume $PWD == "google_research/".
table_path = os.path.join(FLAGS.test_srcdir, testdata_dir,
'tfse-2.tfrecords')
parser_fn = tfse_input_layer.create_tfse_parser(
keypoint_names_2d=(
keypoint_profiles.Std13KeypointProfile2D().keypoint_names),
keypoint_names_3d=(
keypoint_profiles.Std16KeypointProfile3D().keypoint_names),
num_objects=2,
sequence_length=5)
inputs = tfse_input_layer.read_from_table(table_path,
parser_fn=parser_fn)
inputs = next(iter(inputs))
self.assertCountEqual(inputs.keys(), [
'image_sizes', 'keypoints_2d', 'keypoint_scores_2d', 'keypoints_3d'
])
self.assertEqual(inputs['image_sizes'].shape, [2, 5, 2])
self.assertEqual(inputs['keypoints_2d'].shape, [2, 5, 13, 2])
self.assertEqual(inputs['keypoint_scores_2d'].shape, [2, 5, 13])
self.assertEqual(inputs['keypoints_3d'].shape, [2, 5, 16, 3])