Repository: tensorflow/addons

Issue:
Check which image ops are compatible with tf.data.Dataset
Follow up https://github.com/tensorflow/addons/issues/243.

Here is a check list:
+ [x] dense_image_warp
+ [x] sparse_image_warp		
+ [x] interpolate_bilinear	
+ [x] euclidean_distance_transform	
+ [x] adjust_hsv_in_yiq	
+ [x] random_hsv_in_yiq	
+ [x] mean_filter2d	
+ [x] median_filter2d	
+ [x] angles_to_projective_transforms	
+ [x] matrices_to_flat_transforms	
+ [x] rotate	
+ [x] transform

Like what we did in https://github.com/tensorflow/addons/pull/330, feel free to open a pull request for each op:
1. Add test for unknown rank image tensor.
2. Fix it if possible. Otherwise, please notice users that the op is not compatible with unknown rank tensor (shape=None) in its document.



Discussion:
Hello,
Can anyone help me how to test whether a tfa op is compatible with tf.Dataset.
I would like to check for 
- dense_image_warp
- interpolate_bilinear
My question might sound silly, I'm a newbie please help me.
Hi @satishjasthi welcome to addons! So this colab notebook shows what the issue was (has since been fixed for that op):
https://colab.research.google.com/drive/1ZDhnGrorvSf04wzS-0utztZ6Nnzusvl7

Basically, if an image op can't be used in a tf.dataset mapping it's likely because we're not handling an unknown shape correctly. tf.dataset runs as a subgraph so the dynamic shape is not known until the graph is ran. 

A good test to see if it's compatible can be see in previous pull requests that adressed this:
https://github.com/tensorflow/addons/pull/343/files#diff-7137321db62bba0af8193f6f6213fd46R243

Yan added an image util which can be used to infer shape:
https://github.com/tensorflow/addons/blob/master/tensorflow_addons/image/utils.py#L29

Let us know if you run into any issues, and thanks for the help!
regarding the last two ops:

* angles_to_projective_transforms maps a scalar (or batch of N scalars) to an output of dimension (N, 8) representing the equivalent projective transform. It additionally requires image_width and image_height (as scalars or tensors of same dimension as the input angles) to compute the transform.
* matrices_to_flat_transforms similarly maps (N, 3, 3) inputs to (N, 8) outputs (and the inverse flat_transforms_to_matrices from (N, 8) to (N, 3, 3)). For these I don't think there is any sense of image dimensions; only the batch size is not specified.

I don't think there is a sense in which these ops would be suitable for images of unknown dimension as they don't act directly on image data but rather define geometric operations to apply to the image data.

Should we make this issue as done, or is there something to do to make this more transparent to the user?
Hi @mels630, you are right. I tend to close this issue after `sparse_image_warp` gets merged. https://github.com/tensorflow/addons/pull/446. Thanks for pointing it out!
Is there any pending to be added or this issue can be closed?
`sparse_image_warp` was not merged. So check https://github.com/tensorflow/io/issues/839#issuecomment-607362644

Resolving patch:
diff --git a/tensorflow_addons/image/sparse_image_warp.py b/tensorflow_addons/image/sparse_image_warp.py
--- a/tensorflow_addons/image/sparse_image_warp.py
+++ b/tensorflow_addons/image/sparse_image_warp.py
@@ -14,42 +14,48 @@
 # ==============================================================================
 """Image warping using sparse flow defined at control points."""
 
-import numpy as np
 import tensorflow as tf
+
 from tensorflow_addons.image import dense_image_warp
 from tensorflow_addons.image import interpolate_spline
+from tensorflow_addons.image import utils as img_utils
 from tensorflow_addons.utils.types import TensorLike, FloatTensorLike
 
 
 def _get_grid_locations(
     image_height: TensorLike, image_width: TensorLike
 ) -> TensorLike:
-    """Wrapper for `np.meshgrid`."""
-
-    y_range = np.linspace(0, image_height - 1, image_height)
-    x_range = np.linspace(0, image_width - 1, image_width)
-    y_grid, x_grid = np.meshgrid(y_range, x_range, indexing="ij")
-    return np.stack((y_grid, x_grid), -1)
+    """Wrapper for `tf.meshgrid`."""
+    y_range = tf.linspace(0, image_height - 1, image_height)
+    x_range = tf.linspace(0, image_width - 1, image_width)
+    y_grid, x_grid = tf.meshgrid(y_range, x_range, indexing="ij")
+    return tf.stack((y_grid, x_grid), -1)
 
 
-def _expand_to_minibatch(np_array: TensorLike, batch_size: TensorLike) -> TensorLike:
-    """Tile arbitrarily-sized np_array to include new batch dimension."""
-    tiles = [batch_size] + [1] * np_array.ndim
-    return np.tile(np.expand_dims(np_array, 0), tiles)
+def _expand_to_minibatch(array: TensorLike, batch_size: TensorLike) -> TensorLike:
+    """Tile arbitrarily-sized array to include new batch dimension."""
+    batch_size = tf.expand_dims(batch_size, 0)
+    array_ones = tf.ones((tf.rank(array)), dtype=tf.dtypes.int32)
+    tiles = tf.concat([batch_size, array_ones], axis=0)
+    return tf.tile(tf.expand_dims(array, 0), tiles)
 
 
 def _get_boundary_locations(
     image_height: TensorLike, image_width: TensorLike, num_points_per_edge: TensorLike
 ) -> TensorLike:
     """Compute evenly-spaced indices along edge of image."""
-    y_range = np.linspace(0, image_height - 1, num_points_per_edge + 2)
-    x_range = np.linspace(0, image_width - 1, num_points_per_edge + 2)
-    ys, xs = np.meshgrid(y_range, x_range, indexing="ij")
-    is_boundary = np.logical_or(
-        np.logical_or(xs == 0, xs == image_width - 1),
-        np.logical_or(ys == 0, ys == image_height - 1),
+    image_height = tf.cast(image_height, tf.float32)
+    image_width = tf.cast(image_width, tf.float32)
+    y_range = tf.linspace(0.0, image_height - 1, num_points_per_edge + 2)
+    x_range = tf.linspace(0.0, image_width - 1, num_points_per_edge + 2)
+    ys, xs = tf.meshgrid(y_range, x_range, indexing="ij")
+    is_boundary = tf.logical_or(
+        tf.logical_or(tf.equal(xs, 0), tf.equal(xs, image_width - 1)),
+        tf.logical_or(tf.equal(ys, 0), tf.equal(ys, image_height - 1)),
+    )
+    return tf.stack(
+        [tf.boolean_mask(ys, is_boundary), tf.boolean_mask(xs, is_boundary)], axis=-1
     )
-    return np.stack([ys[is_boundary], xs[is_boundary]], axis=-1)
 
 
 def _add_zero_flow_controls_at_boundary(
@@ -85,15 +91,15 @@ def _add_zero_flow_controls_at_boundary(
         image_height, image_width, boundary_points_per_edge
     )
 
-    boundary_point_flows = np.zeros([boundary_point_locations.shape[0], 2])
+    boundary_point_flows = tf.zeros([tf.shape(boundary_point_locations)[0], 2])
 
     type_to_use = control_point_locations.dtype
-    boundary_point_locations = tf.constant(
-        _expand_to_minibatch(boundary_point_locations, batch_size), dtype=type_to_use
+    boundary_point_locations = tf.cast(
+        _expand_to_minibatch(boundary_point_locations, batch_size), type_to_use
     )
 
-    boundary_point_flows = tf.constant(
-        _expand_to_minibatch(boundary_point_flows, batch_size), dtype=type_to_use
+    boundary_point_flows = tf.cast(
+        _expand_to_minibatch(boundary_point_flows, batch_size), type_to_use
     )
 
     merged_control_point_locations = tf.concat(
@@ -140,33 +146,38 @@ def sparse_image_warp(
 
 
     Args:
-      image: `[batch, height, width, channels]` float `Tensor`
-      source_control_point_locations: `[batch, num_control_points, 2]` float
-        `Tensor`
-      dest_control_point_locations: `[batch, num_control_points, 2]` float
-        `Tensor`
+      image: Either a 2-D float `Tensor` of shape `[height, width]`,
+        a 3-D `Tensor` of shape `[height, width, channels]`,
+        or a 4-D `Tensor` of shape `[batch_size, height, width, channels]`.
+        `batch_size` is assumed as one when `image` is a 2-D or 3-D `Tensor`.
+      source_control_point_locations: `[batch_size, num_control_points, 2]` float
+        `Tensor`.
+      dest_control_point_locations: `[batch_size, num_control_points, 2]` float
+        `Tensor`.
       interpolation_order: polynomial order used by the spline interpolation
       regularization_weight: weight on smoothness regularizer in interpolation
       num_boundary_points: How many zero-flow boundary points to include at
         each image edge. Usage:
-          `num_boundary_points=0`: don't add zero-flow points
-          `num_boundary_points=1`: 4 corners of the image
-          `num_boundary_points=2`: 4 corners and one in the middle of each edge
-            (8 points total)
-          `num_boundary_points=n`: 4 corners and n-1 along each edge
+        - `num_boundary_points=0`: don't add zero-flow points
+        - `num_boundary_points=1`: 4 corners of the image
+        - `num_boundary_points=2`: 4 corners and one in the middle of each edge
+          (8 points total)
+        - `num_boundary_points=n`: 4 corners and n-1 along each edge
       name: A name for the operation (optional).
 
-      Note that image and offsets can be of type tf.half, tf.float32, or
-      tf.float64, and do not necessarily have to be the same type.
+      Note that `image` and `offsets` can be of type `tf.half`, `tf.float32`, or
+      `tf.float64`, and do not necessarily have to be the same type.
 
     Returns:
-      warped_image: `[batch, height, width, channels]` float `Tensor` with same
-        type as input image.
-      flow_field: `[batch, height, width, 2]` float `Tensor` containing the
+      warped_image: a float `Tensor` with the same shape and dtype as `image`.
+      flow_field: `[batch_size, height, width, 2]` float `Tensor` containing the
         dense flow field produced by the interpolation.
     """
 
     image = tf.convert_to_tensor(image)
+    original_ndims = img_utils.get_ndims(image)
+    image = img_utils.to_4D_image(image)
+
     source_control_point_locations = tf.convert_to_tensor(
         source_control_point_locations
     )
@@ -178,18 +189,22 @@ def sparse_image_warp(
     boundary_points_per_edge = num_boundary_points - 1
 
     with tf.name_scope(name or "sparse_image_warp"):
-
-        batch_size, image_height, image_width, _ = image.get_shape().as_list()
+        image_shape = tf.shape(image)
+        batch_size, image_height, image_width = (
+            image_shape[0],
+            image_shape[1],
+            image_shape[2],
+        )
 
         # This generates the dense locations where the interpolant
         # will be evaluated.
         grid_locations = _get_grid_locations(image_height, image_width)
 
-        flattened_grid_locations = np.reshape(
+        flattened_grid_locations = tf.reshape(
             grid_locations, [image_height * image_width, 2]
         )
 
-        flattened_grid_locations = tf.constant(
+        flattened_grid_locations = tf.cast(
             _expand_to_minibatch(flattened_grid_locations, batch_size), image.dtype
         )
 
@@ -219,4 +234,4 @@ def sparse_image_warp(
 
         warped_image = dense_image_warp(image, dense_flows)
 
-        return warped_image, dense_flows
+        return img_utils.from_4D_image(warped_image, original_ndims), dense_flows


Test patch:
diff --git a/tensorflow_addons/image/tests/sparse_image_warp_test.py b/tensorflow_addons/image/tests/sparse_image_warp_test.py
--- a/tensorflow_addons/image/tests/sparse_image_warp_test.py
+++ b/tensorflow_addons/image/tests/sparse_image_warp_test.py
@@ -15,7 +15,9 @@
 """Tests for sparse_image_warp."""
 
 import numpy as np
+import pytest
 import tensorflow as tf
+
 from tensorflow_addons.image import sparse_image_warp
 from tensorflow_addons.image.sparse_image_warp import _get_boundary_locations
 from tensorflow_addons.image.sparse_image_warp import _get_grid_locations
@@ -243,3 +245,38 @@ def test_that_backprop_runs():
 
     gradients = t.gradient(warped_image, image).numpy()
     assert np.sum(np.abs(gradients)) != 0
+
+
+@pytest.mark.usefixtures("maybe_run_functions_eagerly")
+@pytest.mark.parametrize("shape", [(9, 12), (9, 12, 3), (1, 9, 12, 3)])
+@pytest.mark.parametrize("interpolation_order", [1, 2, 3])
+@pytest.mark.parametrize("num_boundary_points", [1, 2, 3])
+def test_unknown_shape(shape, interpolation_order, num_boundary_points):
+    control_point_locations = np.asarray([3.0, 3.0]).reshape(1, 1, 2).astype(np.float32)
+    control_point_displacements = (
+        np.asarray([0.25, -0.5]).reshape(1, 1, 2).astype(np.float32)
+    )
+    fn = tf.function(sparse_image_warp).get_concrete_function(
+        image=tf.TensorSpec(shape=None, dtype=tf.float32),
+        source_control_point_locations=tf.TensorSpec(shape=[1, 1, 2], dtype=tf.float32),
+        dest_control_point_locations=tf.TensorSpec(shape=[1, 1, 2], dtype=tf.float32),
+        interpolation_order=interpolation_order,
+        num_boundary_points=num_boundary_points,
+    )
+    image = tf.ones(shape=shape, dtype=tf.float32)
+    expected_output = sparse_image_warp(
+        image,
+        control_point_locations,
+        control_point_locations + control_point_displacements,
+        interpolation_order=interpolation_order,
+        num_boundary_points=num_boundary_points,
+    )
+    output = fn(
+        image,
+        control_point_locations,
+        control_point_locations + control_point_displacements,
+        interpolation_order=interpolation_order,
+        num_boundary_points=num_boundary_points,
+    )
+    np.testing.assert_equal(output[0].numpy(), expected_output[0].numpy())
+    np.testing.assert_equal(output[1].numpy(), expected_output[1].numpy())