Repository: tensorflow/addons

Issue:
Make image transforms compatible with TF Dataset
From #242 we found that the tfa.image transforms are not compatible with dataset mapping. This is a very likely usecase for tfa.image so we need to address this. Here is a minimal example showing the differences in tfa.rotate(pi/2) and tf.image.rot90:
https://colab.research.google.com/drive/1ZDhnGrorvSf04wzS-0utztZ6Nnzusvl7

Problems seem to occur because Dataset mapping runs in graph mode and this check fails:
https://github.com/tensorflow/addons/blob/master/tensorflow_addons/image/transform_ops.py#L323

By comparison here is tf.image.rot90:
https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/ops/image_ops_impl.py#L514



Discussion:
Will think about this more tonight
According to https://github.com/tensorflow/tensorflow/issues/27811, TF2 transforms Data mappings to subgraphs for better performance.

Here is a non-performatic workaround to force Eager Mode and use tfa.image until there is a solution:
image_ds = path_ds.map(lambda path: tf.py_function(func=load_and_preprocess_image,inp=[path],Tout=tf.float32))

I noticed tfa.image.rotate is running on the CPU, so it is running very slowly, is this the expected behaviour?
Thanks for the temporary workaround! Yeah gpu kernels are on the way #118 
Two thoughts:

1.  I opened #244 to have a (failing) test case for this issue
2.  The requirement of fixed and static image rank is not only in rotate, but upstream in `transform`: https://github.com/tensorflow/addons/blob/master/tensorflow_addons/image/transform_ops.py#L76  Is there a strong technical reason why this restriction exists?
3.  The `rot90()` code seems to assume that if no shape information is available, one can assume that the rank is 3.  Do we think that's safe for our transformations as well?
> 1. I opened #244 to have a (failing) test case for this issue

Thank you, @kyleabeauchamp!

> 2. The requirement of fixed and static image rank is not only in rotate, but upstream in `transform`: https://github.com/tensorflow/addons/blob/master/tensorflow_addons/image/transform_ops.py#L76  Is there a strong technical reason why this restriction exists?

I think it is because the cpp implementation [here](https://github.com/tensorflow/addons/blob/master/tensorflow_addons/custom_ops/image/cc/kernels/image_projective_transform_op.cc#L75) requires a static rank of 4 and each dimension will be fetched out for later computation.

> 3. The `rot90()` code seems to assume that if no shape information is available, one can assume that the rank is 3.  Do we think that's safe for our transformations as well?

Well, I do not think it is very safe though. I'd suppose it is designed for `tf.dataset`. (Not so sure about this). Hi, @facaiy, would you mind taking a look at this?
@WindQAQ Thanks for ping me, Tzu-Wei.

> Well, I do not think it is very safe though. I'd suppose it is designed for tf.dataset.

Yeah, it's not safe, and we can relax the restriction if we document it clearly: data must be 4-D.

@mrry Hi, Derek, could you take a look? The rank of tensor seems unknown even in eager mode. 

@karmel Karmel, do you know who could answer those questions about tf.data and tf.image? Thank you.
@jsimsa -- can you advise?
One option is to fix the `addons` logic as discussed in this thread. Another option is to replace `decode_image` with `decode_png`. The problem arises because `decode_image` is executed by tf.data in graph-mode and shape inference fails to determine which one of the decode utilities is used (which seems to trip some `addons` logic).
@jsimsa thank you for the answer and apologies for the slow reply back. I confirmed that switching to `decode_png` does corrrectly infer the shape which is great. One lingering question though:

1) It's true that the addons logic got tripped up by the None shape coming out of `decode_image`... but in `tf.image.rot90` None shape is handled by defaulting to a 3D shape. Is there a reason this is safe to do? We can copy the same logic but I'm not sure I'm following the rationale.


I tried digging through git history to find the person who might have the answer for you and as far as I could tell is should be @aselle who introduced that logic in 2016.
Hi @aselle Do you happen to have any comment on why its safe to default to 3D rotate when shape is `None`?
https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/ops/image_ops_impl.py#L511
It probably is just so that if you don't know anything about the shape you can still proceed.  None of ndims means the rank is not known. so it's just a fallback. 
So we should handle unknown rank, because it's valid in graph mode. And it seems that there are at least two solutions:
1. use dynamic shape if possible;
2. fallback to 3D data (perhaps unsafe).

Resolving patch:
diff --git a/tensorflow_addons/image/transform_ops.py b/tensorflow_addons/image/transform_ops.py
--- a/tensorflow_addons/image/transform_ops.py
+++ b/tensorflow_addons/image/transform_ops.py
@@ -18,6 +18,7 @@
 from __future__ import print_function
 
 import tensorflow as tf
+from tensorflow_addons.image import utils as img_utils
 from tensorflow_addons.utils.resource_loader import get_path_to_datafile
 
 _image_ops_so = tf.load_op_library(
@@ -40,8 +41,7 @@ def transform(images,
     Args:
       images: A tensor of shape (num_images, num_rows, num_columns,
         num_channels) (NHWC), (num_rows, num_columns, num_channels) (HWC), or
-        (num_rows, num_columns) (HW). The rank must be statically known (the
-        shape is not `TensorShape(None)`.
+        (num_rows, num_columns) (HW).
       transforms: Projective transform matrix/matrices. A vector of length 8 or
         tensor of size N x 8. If one row of transforms is
         [a0, a1, a2, b0, b1, b2, c0, c1], then it maps the *output* point
@@ -72,16 +72,8 @@ def transform(images,
             transforms, name="transforms", dtype=tf.dtypes.float32)
         if image_or_images.dtype.base_dtype not in _IMAGE_DTYPES:
             raise TypeError("Invalid dtype %s." % image_or_images.dtype)
-        elif image_or_images.get_shape().ndims is None:
-            raise TypeError("image_or_images rank must be statically known")
-        elif len(image_or_images.get_shape()) == 2:
-            images = image_or_images[None, :, :, None]
-        elif len(image_or_images.get_shape()) == 3:
-            images = image_or_images[None, :, :, :]
-        elif len(image_or_images.get_shape()) == 4:
-            images = image_or_images
-        else:
-            raise TypeError("Images should have rank between 2 and 4.")
+        images = img_utils.to_4D_image(image_or_images)
+        original_ndims = img_utils.get_ndims(image_or_images)
 
         if output_shape is None:
             output_shape = tf.shape(images)[1:3]
@@ -109,12 +101,7 @@ def transform(images,
             output_shape=output_shape,
             transforms=transforms,
             interpolation=interpolation.upper())
-        if len(image_or_images.get_shape()) == 2:
-            return output[0, :, :, 0]
-        elif len(image_or_images.get_shape()) == 3:
-            return output[0, :, :, :]
-        else:
-            return output
+        return img_utils.from_4D_image(output, original_ndims)
 
 
 @tf.function
@@ -299,8 +286,7 @@ def rotate(images, angles, interpolation="NEAREST", name=None):
       images: A tensor of shape
         (num_images, num_rows, num_columns, num_channels)
         (NHWC), (num_rows, num_columns, num_channels) (HWC), or
-        (num_rows, num_columns) (HW). The rank must be statically known (the
-        shape is not `TensorShape(None)`.
+        (num_rows, num_columns) (HW).
       angles: A scalar angle to rotate all images by, or (if images has rank 4)
         a vector of length num_images, with an angle for each image in the
         batch.
@@ -319,16 +305,8 @@ def rotate(images, angles, interpolation="NEAREST", name=None):
         image_or_images = tf.convert_to_tensor(images)
         if image_or_images.dtype.base_dtype not in _IMAGE_DTYPES:
             raise TypeError("Invalid dtype %s." % image_or_images.dtype)
-        if image_or_images.get_shape().ndims is None:
-            raise TypeError("image_or_images rank must be statically known")
-        elif len(image_or_images.get_shape()) == 2:
-            images = image_or_images[None, :, :, None]
-        elif len(image_or_images.get_shape()) == 3:
-            images = image_or_images[None, :, :, :]
-        elif len(image_or_images.get_shape()) == 4:
-            images = image_or_images
-        else:
-            raise TypeError("Images should have rank between 2 and 4.")
+        images = img_utils.to_4D_image(image_or_images)
+        original_ndims = img_utils.get_ndims(image_or_images)
 
         image_height = tf.cast(tf.shape(images)[1], tf.dtypes.float32)[None]
         image_width = tf.cast(tf.shape(images)[2], tf.dtypes.float32)[None]
@@ -336,11 +314,4 @@ def rotate(images, angles, interpolation="NEAREST", name=None):
             images,
             angles_to_projective_transforms(angles, image_height, image_width),
             interpolation=interpolation)
-        if image_or_images.get_shape().ndims is None:
-            raise TypeError("image_or_images rank must be statically known")
-        elif len(image_or_images.get_shape()) == 2:
-            return output[0, :, :, 0]
-        elif len(image_or_images.get_shape()) == 3:
-            return output[0, :, :, :]
-        else:
-            return output
+        return img_utils.from_4D_image(output, original_ndims)
diff --git a/tensorflow_addons/image/utils.py b/tensorflow_addons/image/utils.py
new file mode 100644
--- /dev/null
+++ b/tensorflow_addons/image/utils.py
@@ -0,0 +1,108 @@
+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Image util ops."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import tensorflow as tf
+
+
+def get_ndims(image):
+    return image.get_shape().ndims or tf.rank(image)
+
+
+@tf.function
+def to_4D_image(image):
+    """Convert 2/3/4D image to 4D image.
+
+    Args:
+      image: 2/3/4D tensor.
+
+    Returns:
+      4D tensor with the same type.
+    """
+    # yapf:disable
+    with tf.control_dependencies([
+        tf.debugging.assert_rank_in(
+            image, [2, 3, 4], message='`image` must be 2/3/4D tensor')
+    ]):
+        # yapf: enable
+        ndims = image.get_shape().ndims
+        if ndims is None:
+            return _dynamic_to_4D_image(image)
+        elif ndims == 2:
+            return image[None, :, :, None]
+        elif ndims == 3:
+            return image[None, :, :, :]
+        else:
+            return image
+
+
+def _dynamic_to_4D_image(image):
+    shape = tf.shape(image)
+    original_rank = tf.rank(image)
+    # 4D image => [N, H, W, C] or [N, C, H, W]
+    # 3D image => [1, H, W, C] or [1, C, H, W]
+    # 2D image => [1, H, W, 1]
+    left_pad = tf.cast(tf.less_equal(original_rank, 3), dtype=tf.int32)
+    right_pad = tf.cast(tf.equal(original_rank, 2), dtype=tf.int32)
+    # yapf: disable
+    new_shape = tf.concat(
+        [tf.ones(shape=left_pad, dtype=tf.int32),
+         shape,
+         tf.ones(shape=right_pad, dtype=tf.int32)],
+        axis=0)
+    # yapf: enable
+    return tf.reshape(image, new_shape)
+
+
+@tf.function
+def from_4D_image(image, ndims):
+    """Convert back to an image with `ndims` rank.
+
+    Args:
+      image: 4D tensor.
+      ndims: The original rank of the image.
+
+    Returns:
+      `ndims`-D tensor with the same type.
+    """
+    # yapf:disable
+    with tf.control_dependencies([
+        tf.debugging.assert_rank(
+            image, 4, message='`image` must be 4D tensor')
+    ]):
+        # yapf:enable
+        if isinstance(ndims, tf.Tensor):
+            return _dynamic_from_4D_image(image, ndims)
+        elif ndims == 2:
+            return tf.squeeze(image, [0, 3])
+        elif ndims == 3:
+            return tf.squeeze(image, [0])
+        else:
+            return image
+
+
+def _dynamic_from_4D_image(image, original_rank):
+    shape = tf.shape(image)
+    # 4D image <= [N, H, W, C] or [N, C, H, W]
+    # 3D image <= [1, H, W, C] or [1, C, H, W]
+    # 2D image <= [1, H, W, 1]
+    begin = tf.cast(tf.less_equal(original_rank, 3), dtype=tf.int32)
+    end = 4 - tf.cast(tf.equal(original_rank, 2), dtype=tf.int32)
+    new_shape = shape[begin:end]
+    return tf.reshape(image, new_shape)


Test patch:
diff --git a/tensorflow_addons/image/transform_ops_test.py b/tensorflow_addons/image/transform_ops_test.py
--- a/tensorflow_addons/image/transform_ops_test.py
+++ b/tensorflow_addons/image/transform_ops_test.py
@@ -72,6 +72,15 @@ def test_transform_static_output_shape(self):
             output_shape=tf.constant([3, 5]))
         self.assertAllEqual([3, 5], result.shape)
 
+    @test_utils.run_in_graph_and_eager_modes
+    def test_transform_unknown_shape(self):
+        fn = transform_ops.transform.get_concrete_function(
+            tf.TensorSpec(shape=None, dtype=tf.float32),
+            [1, 0, 0, 0, 1, 0, 0, 0])
+        for shape in (2, 4), (2, 4, 3), (1, 2, 4, 3):
+            image = tf.ones(shape=shape)
+            self.assertAllEqual(self.evaluate(image), self.evaluate(fn(image)))
+
     def _test_grad(self, input_shape, output_shape=None):
         image_size = tf.math.cumprod(input_shape)[-1]
         image_size = tf.cast(image_size, tf.float32)
@@ -270,6 +279,13 @@ def test_rotate_static_shape(self):
             image, tf.random.uniform((), -1, 1), interpolation="BILINEAR")
         self.assertEqual(image.get_shape(), result.get_shape())
 
+    def test_unknown_shape(self):
+        fn = transform_ops.rotate.get_concrete_function(
+            tf.TensorSpec(shape=None, dtype=tf.float32), 0)
+        for shape in (2, 4), (2, 4, 3), (1, 2, 4, 3):
+            image = tf.ones(shape=shape)
+            self.assertAllEqual(self.evaluate(image), self.evaluate(fn(image)))
+
 
 if __name__ == "__main__":
     tf.test.main()
diff --git a/tensorflow_addons/image/utils_test.py b/tensorflow_addons/image/utils_test.py
new file mode 100644
--- /dev/null
+++ b/tensorflow_addons/image/utils_test.py
@@ -0,0 +1,99 @@
+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for util ops."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import tensorflow as tf
+
+from tensorflow_addons.image import utils as img_utils
+from tensorflow_addons.utils import test_utils
+
+
+@test_utils.run_all_in_graph_and_eager_modes
+class UtilsOpsTest(tf.test.TestCase):
+    def test_to_4D_image(self):
+        for shape in (2, 4), (2, 4, 1), (1, 2, 4, 1):
+            exp = tf.ones(shape=(1, 2, 4, 1))
+            res = img_utils.to_4D_image(tf.ones(shape=shape))
+            # static shape:
+            self.assertAllEqual(exp.get_shape(), res.get_shape())
+            self.assertAllEqual(self.evaluate(exp), self.evaluate(res))
+
+    def test_to_4D_image_with_unknown_shape(self):
+        fn = img_utils.to_4D_image.get_concrete_function(
+            tf.TensorSpec(shape=None, dtype=tf.float32))
+        for shape in (2, 4), (2, 4, 1), (1, 2, 4, 1):
+            exp = tf.ones(shape=(1, 2, 4, 1))
+            res = fn(tf.ones(shape=shape))
+            self.assertAllEqual(self.evaluate(exp), self.evaluate(res))
+
+    def test_to_4D_image_with_invalid_shape(self):
+        errors = (ValueError, tf.errors.InvalidArgumentError)
+        with self.assertRaisesRegexp(errors, '`image` must be 2/3/4D tensor'):
+            img_utils.to_4D_image(tf.ones(shape=(1,)))
+
+        with self.assertRaisesRegexp(errors, '`image` must be 2/3/4D tensor'):
+            img_utils.to_4D_image(tf.ones(shape=(1, 2, 4, 3, 2)))
+
+    def test_from_4D_image(self):
+        for shape in (2, 4), (2, 4, 1), (1, 2, 4, 1):
+            exp = tf.ones(shape=shape)
+            res = img_utils.from_4D_image(
+                tf.ones(shape=(1, 2, 4, 1)), len(shape))
+            # static shape:
+            self.assertAllEqual(exp.get_shape(), res.get_shape())
+            self.assertAllEqual(self.evaluate(exp), self.evaluate(res))
+
+    def test_from_4D_image_with_unknown_shape(self):
+        for shape in (2, 4), (2, 4, 1), (1, 2, 4, 1):
+            exp = tf.ones(shape=shape)
+            fn = img_utils.from_4D_image.get_concrete_function(
+                tf.TensorSpec(shape=None, dtype=tf.float32), tf.size(shape))
+            res = fn(tf.ones(shape=(1, 2, 4, 1)), tf.size(shape))
+            self.assertAllEqual(self.evaluate(exp), self.evaluate(res))
+
+    def test_from_4D_image_with_invalid_data(self):
+        with self.assertRaises(ValueError):
+            self.evaluate(
+                img_utils.from_4D_image(tf.ones(shape=(2, 2, 4, 1)), 2))
+
+        with self.assertRaises(tf.errors.InvalidArgumentError):
+            self.evaluate(
+                img_utils.from_4D_image(
+                    tf.ones(shape=(2, 2, 4, 1)), tf.constant(2)))
+
+    def test_from_4D_image_with_invalid_shape(self):
+        errors = (ValueError, tf.errors.InvalidArgumentError)
+        for rank in 2, tf.constant(2):
+            with self.subTest(rank=rank):
+                with self.assertRaisesRegexp(errors,
+                                             '`image` must be 4D tensor'):
+                    img_utils.from_4D_image(tf.ones(shape=(2, 4)), rank)
+
+                with self.assertRaisesRegexp(errors,
+                                             '`image` must be 4D tensor'):
+                    img_utils.from_4D_image(tf.ones(shape=(2, 4, 1)), rank)
+
+                with self.assertRaisesRegexp(errors,
+                                             '`image` must be 4D tensor'):
+                    img_utils.from_4D_image(
+                        tf.ones(shape=(1, 2, 4, 1, 1)), rank)
+
+
+if __name__ == "__main__":
+    tf.test.main()