tf.image.flip_left_right and tf.image.flip_up_down incorrectly assumes rank-3 image and flips along the wrong axis

[palotasb-conti]

System information

• Have I written custom code (as opposed to using a stock example script provided in TensorFlow): Yes, see below
• OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Windows 10
• TensorFlow installed from (source or binary): Binary
• TensorFlow version (use command below): 2.0
• Python version: 3.7

You can also obtain the TensorFlow version with:
2. TF 2.0: python -c "import tensorflow as tf; print(tf.version.GIT_VERSION, tf.version.VERSION)"
unknown 2.0.0

Describe the current behavior

Setup: Under some circumstances (in my case, when a tf.Tensor is inside a tf.data.Dataset and it is being conditionally transformed inside a tf.data.Dataset.map() call), the shape of a Tensor becomes unknown (this may be expected behavior still).

Behavior: When tf.image.flip_left_right is applied to this tf.Tensor, the function incorrectly assumes a rank-3 shape and flips the image along the wrong axis.

Describe the expected behavior

When tf.image.flip_left_right is applied to the tf.Tensor, the function does not assume a rank-3 shape and flips the image along the correct axis.

Standalone code to reproduce the issue

import tensorflow as tf

def correct_image_flip_left_right(image):
    return tf.reverse(image, axis=[-2])

PATCH_TF = False  # Change this to True to fix the bug
if PATCH_TF:
    tf.image.flip_left_right = correct_image_flip_left_right

image_input = tf.convert_to_tensor([
    # batch: 0
    [
        # y: 0
        [[0, 1, 2], [3, 4, 5]],  # x=0,1 channels=0,1,2
        # y: 1
        [[6, 7, 8], [9, 10, 11]],  # x=0,1 channels=0,1,2
    ],
])

image_flipped_directly = tf.image.flip_left_right(image_input)

expected_output = tf.convert_to_tensor([
    # batch: 0
    [
        # y: 0
        [[3, 4, 5], [0, 1, 2]],  # x=0,1 channels=0,1,2
        # y: 1
        [[9, 10, 11], [6, 7, 8]],  # x=0,1 channels=0,1,2
    ],
])

tf.assert_equal(image_flipped_directly, expected_output)
print("Directly calling tf.image.flip_left_right works as exected.")

def generator():  yield image_input

dataset = tf.data.Dataset.from_generator(generator, output_types=tf.int32)

def flip_it(image, do_flip: bool):
    if do_flip:
        return tf.image.flip_left_right(image)
    else:
        return image

dataset = dataset.map(lambda image: flip_it(image, tf.constant(True)))

image_flipped_via_dataset_map = next(iter(dataset))

print("image_flipped_via_dataset_map:")
print(image_flipped_via_dataset_map)
print("expected_output:")
print(expected_output)

tf.assert_equal(image_flipped_via_dataset_map, expected_output)
# This assertion fails even though it shouldn't unless PATCH_TF is True

print("If you can see this message, image_flipped_via_dataset_map == expected_output")

Output:

> python .\tf_image_flip.py
2020-06-18 17:10:08.875275: I tensorflow/core/platform/cpu_feature_guard.cc:142] Your CPU supports instructions that this TensorFlow binary was not compiled to use: AVX AVX2
Directly calling tf.image.flip_left_right works as exected.
image_flipped_via_dataset_map:
tf.Tensor(
[[[[ 6  7  8]
   [ 9 10 11]]

  [[ 0  1  2]
   [ 3  4  5]]]], shape=(1, 2, 2, 3), dtype=int32)
expected_output:
tf.Tensor(
[[[[ 3  4  5]
   [ 0  1  2]]

  [[ 9 10 11]
   [ 6  7  8]]]], shape=(1, 2, 2, 3), dtype=int32)
Traceback (most recent call last):
  File ".\tf_image_flip.py", line 54, in <module>
    tf.assert_equal(image_flipped_via_dataset_map, expected_output)
  File "C:\Users\uib58003\AppData\Local\Continuum\miniconda3\lib\site-packages\tensorflow_core\python\ops\check_ops.py", line 456, in assert_equal_v2
    return assert_equal(x=x, y=y, summarize=summarize, message=message, name=name)
  File "C:\Users\uib58003\AppData\Local\Continuum\miniconda3\lib\site-packages\tensorflow_core\python\ops\check_ops.py", line 546, in assert_equal
    (message or '', index_and_values_str, summary_msg)))
tensorflow.python.framework.errors_impl.InvalidArgumentError:
Condition x == y did not hold.
Indices of first 3 different values:
[[0 0 0 0]
 [0 0 0 1]
 [0 0 0 2]]
Corresponding x values:
[6 7 8]
Corresponding y values:
[3 4 5]
First 3 elements of x:
[6 7 8]
First 3 elements of y:
[3 4 5]

Other info / logs

The bug is in tensorflow_core.ops.image_ops_impl._flip, called by image_ops_impl.flip_left_right:

tensorflow/tensorflow/python/ops/image_ops_impl.py

Lines 537 to 546 in 2b96f36

	with ops.name_scope(None, scope_name, [image]):
	image = ops.convert_to_tensor(image, name='image')
	image = _AssertAtLeast3DImage(image)
	shape = image.get_shape()
	if shape.ndims == 3 or shape.ndims is None:
	return fix_image_flip_shape(image, array_ops.reverse(image, [flip_index]))
	elif shape.ndims == 4:
	return array_ops.reverse(image, [flip_index + 1])
	else:
	raise ValueError('\'image\' must have either 3 or 4 dimensions.')

Tensorflow should not assume based on shape.ndims is None that the image is 3-dimensional and it should not call fix_image_flip_shape which further builds on this incorrect assumption:

tensorflow/tensorflow/python/ops/image_ops_impl.py

Lines 315 to 320 in 2b96f36

	image_shape = image.get_shape()
	if image_shape == tensor_shape.unknown_shape():
	result.set_shape([None, None, None])
	else:
	result.set_shape(image_shape)
	return result

This causes the flipping to happen along the wrong axis.

The flip_left_right function should be implemented by array_ops.reverse(image, axis=[-2]) which always flips along the correct axis for any number of ranks as long as the dimensions end in ..., HEIGHT, WIDTH, CHANNELS].

The same bug probably appears in tf.image.flip_up_down for the same reason based on looking at the source code, but I haven't tested this. That function should always apply tf.reverse(image, axis=[-3]) for the same reason as before.

[ravikyram]

I have tried in colab with TF 2.0,2.1,2.2, nightly versions and was able to reproduce the issue.Please, find the gist here.Thanks!

[yongtang]

Added a PR #40626 for the fix.

[google-ml-butler]

Are you satisfied with the resolution of your issue?
Yes
No