import tensorflow as tf

def _fake_image_resizer_fn(image, mask):
  return (image, mask, tf.shape(image))


class DataTransformationFnTest(test_case.TestCase):

  def test_combine_additional_channels_if_present(self):
    image = np.random.rand(4, 4, 3).astype(np.float32)
    additional_channels = np.random.rand(4, 4, 2).astype(np.float32)
    tensor_dict = {
        fields.InputDataFields.image:
            tf.constant(image),
        fields.InputDataFields.image_additional_channels:
            tf.constant(additional_channels),
        fields.InputDataFields.groundtruth_classes:
            tf.constant(np.array([1, 1], np.int32))
    }

    input_transformation_fn = functools.partial(
        inputs.transform_input_data,
        model_preprocess_fn=_fake_model_preprocessor_fn,
        image_resizer_fn=_fake_image_resizer_fn,
