import tensorflow as tf


    self.assertAllEqual(augmented_tensor_dict[fields.InputDataFields.image],
                        np_image + 1)
    self.assertAllEqual(
        augmented_tensor_dict[fields.InputDataFields.groundtruth_classes],
        [[0, 0, 0, 1], [0, 1, 0, 0]])

  def test_applies_data_augmentation_fn_before_model_preprocess_fn(self):
    np_image = np.random.randint(256, size=(4, 4, 3))
    tensor_dict = {
        fields.InputDataFields.image:
            tf.constant(np_image),
        fields.InputDataFields.groundtruth_classes:
            tf.constant(np.array([3, 1], np.int32))
    }

    def mul_two_model_preprocessor_fn(image):
      return (image * 2, tf.expand_dims(tf.shape(image)[1:], axis=0))

    def add_five_to_image_data_augmentation_fn(tensor_dict):
      tensor_dict[fields.InputDataFields.image] += 5
      return tensor_dict
