def test_get_valid_keypoints_mask(self):

    def graph_fn():
        class_onehot = tf.constant([[0, 0, 1, 0, 0], [0, 1, 0, 0, 0], [0, 0, 1, 0, 1]], dtype=tf.float32)
        keypoints = tf.constant([[0.1, float('nan'), 0.2, 0.0], [0.0, 0.0, 0.1, 0.9], [3.2, 4.3, float('nan'), 0.2]], dtype=tf.float32)
        keypoint_coordinates = tf.stack([keypoints, keypoints], axis=2)
        (mask, keypoints_nan_to_zeros) = ta_utils.get_valid_keypoint_mask_for_class(keypoint_coordinates=keypoint_coordinates, class_id=2, class_onehot=class_onehot, keypoint_indices=[1, 2])
        return (mask, keypoints_nan_to_zeros)
    keypoints = np.array([[0.0, 0.2], [0.0, 0.1], [4.3, 0.0]])
    expected_mask = np.array([[0, 1], [0, 0], [1, 0]])
    expected_keypoints = np.stack([keypoints, keypoints], axis=2)
    (mask, keypoints_nan_to_zeros) = self.execute(graph_fn, [])
    np.testing.assert_array_equal(mask, expected_mask)
    np.testing.assert_array_almost_equal(keypoints_nan_to_zeros, expected_keypoints)
def test_get_valid_keypoints_mask(self):

    def graph_fn():
        class_onehot = tf.constant([[0, 0, 1, 0, 0], [0, 1, 0, 0, 0], [0, 0, 1, 0, 1]], dtype=tf.float32)
        keypoints = tf.constant([[0.1, float('nan'), 0.2, 0.0], [0.0, 0.0, 0.1, 0.9], [3.2, 4.3, float('nan'), 0.2]], dtype=tf.float32)
        keypoint_coordinates = tf.stack([keypoints, keypoints], axis=2)
        (mask, keypoints_nan_to_zeros) = ta_utils.get_valid_keypoint_mask_for_class(keypoint_coordinates=keypoint_coordinates, class_id=2, class_onehot=class_onehot, keypoint_indices=[1, 2])
        return (mask, keypoints_nan_to_zeros)
    keypoints = np.array([[0.0, 0.2], [0.0, 0.1], [4.3, 0.0]])
    expected_mask = np.array([[0, 1], [0, 0], [1, 0]])
    expected_keypoints = np.stack([keypoints, keypoints], axis=2)
    (mask, keypoints_nan_to_zeros) = self.execute(graph_fn, [])
    np.testing.assert_array_equal(mask, expected_mask)
    np.testing.assert_array_almost_equal(keypoints_nan_to_zeros, expected_keypoints)