def test_insert_zeros_at_indices(self):
keypoints = tf.constant([[[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]]])
indices = [1, 1, 3]
keypoints = keypoint_utils.insert_at_indices(
keypoints, indices, insert_keypoints=None)
self.assertAllClose(keypoints, [[[1.0, 2.0], [0.0, 0.0], [0.0, 0.0],
[3.0, 4.0], [5.0, 6.0], [0.0, 0.0]]])
def test_insert_at_indices(self):
keypoints = tf.constant([[[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]]])
indices = [1, 3, 3]
insert_keypoints = tf.constant([[[7.0, 8.0], [9.0, 10.0], [11.0, 12.0]]])
keypoints = keypoint_utils.insert_at_indices(
keypoints, indices, insert_keypoints=insert_keypoints)
self.assertAllClose(keypoints, [[[1.0, 2.0], [7.0, 8.0], [3.0, 4.0],
[5.0, 6.0], [9.0, 10.0], [11.0, 12.0]]])
