Repository: tensorflow/addons

Issue:
tfa.image.translate_xy, tfa.image.shear_* fails to work with float images
**System information**
- OS Platform and Distribution (e.g., Linux Ubuntu 16.04): `Google Colaboratory` (Debian) / `macOS v10.15.6`
- TensorFlow version and how it was installed (source or binary): `Binary TF_v2.3.0`
- TensorFlow-Addons version and how it was installed (source or binary): `Binary tfa-nightly_v0.12.0.dev20200811022730`
- Python version: `3.7.6`
- Is GPU used? (yes/no): `no`

**Describe the bug**

`tfa.image.shear_x`, `tfa.image.shear_y`, `tfa.image.translate_xy` and possibly other args that make use of the `replace` argument fails to work with float images. The `replace` argument should strictly be an integer (uint8) which is also suggested in the documentation, but the same results in either `TypeError` or `ConcatV2` errors when working with non-integer images.

**Code to reproduce the issue**

```python3
float_image = tf.random.uniform([32, 32, 3], 0, 1, dtype=tf.float32)
sheared_image = tfa.image.shear_x(float_image, 0.3, replace=0.5)  # fails with TypeError
sheared_image = tfa.image.shear_x(float_image, 0.3, replace=128)  # fails with Concat error
```

A more detailed Colab with the stack trace and other possible combinations that still won't work for the same reason can be found in [[tfa.image] op bug w/ replace argument.ipynb](https://colab.research.google.com/gist/swghosh/7245f498ba407f5552a304fa2784c0ac/-tfa-image-op-bug-w-replace-argument.ipynb)

**Other info / logs**

```python3
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-12-50c68fed8a6d> in <module>()
----> 1 sheared_image3 = tfa.image.shear_x(float_image, 0.3, replace=0.5)
      2 sheared_image4 = tfa.image.shear_y(float_image, 0.4, replace=0.4)
      3 
      4 plt.imshow(sheared_image3)
      5 plt.show()

5 frames
/usr/local/lib/python3.6/dist-packages/tensorflow_addons/image/transform_ops.py in shear_x(image, level, replace)
    324     #  0  1].
    325     image = transform(wrap(image), [1.0, level, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0])
--> 326     return unwrap(image, replace)
    327 
    328 

/usr/local/lib/python3.6/dist-packages/tensorflow_addons/image/utils.py in unwrap(image, replace)
    135     alpha_channel = flattened_image[:, 3]
    136 
--> 137     replace = tf.constant(replace, tf.uint8)
    138     if tf.rank(replace) == 0:
    139         replace = tf.expand_dims(replace, 0)

/usr/local/lib/python3.6/dist-packages/tensorflow/python/framework/constant_op.py in constant(value, dtype, shape, name)
    262   """
    263   return _constant_impl(value, dtype, shape, name, verify_shape=False,
--> 264                         allow_broadcast=True)
    265 
    266 

/usr/local/lib/python3.6/dist-packages/tensorflow/python/framework/constant_op.py in _constant_impl(value, dtype, shape, name, verify_shape, allow_broadcast)
    273       with trace.Trace("tf.constant"):
    274         return _constant_eager_impl(ctx, value, dtype, shape, verify_shape)
--> 275     return _constant_eager_impl(ctx, value, dtype, shape, verify_shape)
    276 
    277   g = ops.get_default_graph()

/usr/local/lib/python3.6/dist-packages/tensorflow/python/framework/constant_op.py in _constant_eager_impl(ctx, value, dtype, shape, verify_shape)
    298 def _constant_eager_impl(ctx, value, dtype, shape, verify_shape):
    299   """Implementation of eager constant."""
--> 300   t = convert_to_eager_tensor(value, ctx, dtype)
    301   if shape is None:
    302     return t

/usr/local/lib/python3.6/dist-packages/tensorflow/python/framework/constant_op.py in convert_to_eager_tensor(value, ctx, dtype)
     96       dtype = dtypes.as_dtype(dtype).as_datatype_enum
     97   ctx.ensure_initialized()
---> 98   return ops.EagerTensor(value, ctx.device_name, dtype)
     99 
    100 

TypeError: Cannot convert 0.5 to EagerTensor of dtype uint8
```

```python3
---------------------------------------------------------------------------
InvalidArgumentError                      Traceback (most recent call last)
<ipython-input-13-ffd69f46b291> in <module>()
----> 1 sheared_image3 = tfa.image.shear_x(float_image, 0.3, replace=128)
      2 sheared_image4 = tfa.image.shear_y(float_image, 0.4, replace=0)
      3 
      4 plt.imshow(sheared_image3)
      5 plt.show()

6 frames
/usr/local/lib/python3.6/dist-packages/six.py in raise_from(value, from_value)

InvalidArgumentError: cannot compute ConcatV2 as input #1(zero-based) was expected to be a uint8 tensor but is a float tensor [Op:ConcatV2] name: concat
```

What would be the way out to run these op(s) with float images?

@windqaq @facaiy @sayoojbk
/cc: @tanzhenyu


Discussion:
This is the bug in our part. I can fix it recently. But personally, I discourage to use this function after https://github.com/tensorflow/tensorflow/pull/41365 gets merged. Once one can control `fill_value` when doing transformation, `wrap`/`unwrap` is redundant.
Hi @WindQAQ, thanks for sharing the update on TF `ImageProjectiveTransformV2` and it's upcoming support for `fill_mode` as a complete op arg.
So as a workaround we should wait for the changes to get merged on tf/master and then look into a fix for this (float support on image ops with `replace`), right?
> Hi @WindQAQ, thanks for sharing the update on TF `ImageProjectiveTransformV2` and it's upcoming support for `fill_mode` as a complete op arg.
> So as a workaround we should wait for the changes to get merged on tf/master and then look into a fix for this (float support on image ops with `replace`), right?

Yes, I can fix it today or you can open an PR for it!
>> Hi @WindQAQ, thanks for sharing the update on TF ImageProjectiveTransformV2 and it's upcoming support for fill_mode as a complete op arg.
> So as a workaround we should wait for the changes to get merged on tf/master and then look into a fix for this (float support on image ops with replace), right?

> Yes, I can fix it today or you can open an PR for it!

I'll work on a PR and get this done. (utilising the fill_mode argument from ImageProjectiveTransformV2)
Should I then skip wrap / unwrap completely?
> > > Hi @WindQAQ, thanks for sharing the update on TF ImageProjectiveTransformV2 and it's upcoming support for fill_mode as a complete op arg.
> > > So as a workaround we should wait for the changes to get merged on tf/master and then look into a fix for this (float support on image ops with replace), right?
> 
> > Yes, I can fix it today or you can open an PR for it!
> 
> I'll work on a PR and get this done. (utilising the fill_mode argument from ImageProjectiveTransformV2)
> Should I then skip wrap / unwrap completely?

Sorry that I accidentally edited your comment... And just leave it alone as `fill_value` has not supported yet.
Oh cool cool. We can wait till the support is added. Thanks.

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
@@ -308,11 +308,11 @@ def rotate(
         return img_utils.from_4D_image(output, original_ndims)
 
 
-def shear_x(image: TensorLike, level: float, replace: int) -> TensorLike:
+def shear_x(image: TensorLike, level: float, replace: TensorLike) -> TensorLike:
     """Perform shear operation on an image (x-axis).
 
     Args:
-        image: A 3D image Tensor.
+        image: A 3D image `Tensor`.
         level: A float denoting shear element along y-axis
         replace: A one or three value 1D tensor to fill empty pixels.
     Returns:
@@ -327,7 +327,7 @@ def shear_x(image: TensorLike, level: float, replace: int) -> TensorLike:
     return unwrap(image, replace)
 
 
-def shear_y(image: TensorLike, level: float, replace: int) -> TensorLike:
+def shear_y(image: TensorLike, level: float, replace: TensorLike) -> TensorLike:
     """Perform shear operation on an image (y-axis).
 
     Args:
diff --git a/tensorflow_addons/image/translate_ops.py b/tensorflow_addons/image/translate_ops.py
--- a/tensorflow_addons/image/translate_ops.py
+++ b/tensorflow_addons/image/translate_ops.py
@@ -107,13 +107,13 @@ def translate(
 
 
 def translate_xy(
-    image: TensorLike, translate_to: TensorLike, replace: int
+    image: TensorLike, translate_to: TensorLike, replace: TensorLike
 ) -> TensorLike:
     """Translates image in X or Y dimension.
 
     Args:
         image: A 3D image `Tensor`.
-        translate_to: A 1D `Tensor` to translate [x, y]
+        translate_to: A 1D `Tensor` to translate `[x, y]`.
         replace: A one or three value 1D `Tensor` to fill empty pixels.
     Returns:
         Translated image along X or Y axis, with space outside image
diff --git a/tensorflow_addons/image/utils.py b/tensorflow_addons/image/utils.py
--- a/tensorflow_addons/image/utils.py
+++ b/tensorflow_addons/image/utils.py
@@ -134,11 +134,11 @@ def unwrap(image, replace):
     # Find all pixels where the last channel is zero.
     alpha_channel = flattened_image[:, 3]
 
-    replace = tf.constant(replace, tf.uint8)
+    replace = tf.cast(replace, image.dtype)
     if tf.rank(replace) == 0:
         replace = tf.expand_dims(replace, 0)
         replace = tf.concat([replace, replace, replace], 0)
-    replace = tf.concat([replace, tf.ones([1], dtype=image.dtype)], 0)
+    replace = tf.concat([replace, tf.ones([1], dtype=replace.dtype)], 0)
 
     # Where they are zero, fill them in with 'replace'.
     cond = tf.equal(alpha_channel, 1)


Test patch:
diff --git a/tensorflow_addons/image/tests/transform_ops_test.py b/tensorflow_addons/image/tests/transform_ops_test.py
--- a/tensorflow_addons/image/tests/transform_ops_test.py
+++ b/tensorflow_addons/image/tests/transform_ops_test.py
@@ -18,9 +18,11 @@
 import numpy as np
 import tensorflow as tf
 
-from tensorflow_addons.image import transform_ops
 from skimage import transform
 
+from tensorflow_addons.image import transform_ops
+from tensorflow_addons.utils import test_utils
+
 _DTYPES = {
     tf.dtypes.uint8,
     tf.dtypes.int32,
@@ -322,11 +324,13 @@ def test_unknown_shape():
         np.testing.assert_equal(image.numpy(), fn(image).numpy())
 
 
-# TODO: Parameterize on dtypes
 @pytest.mark.usefixtures("maybe_run_functions_eagerly")
-def test_shear_x():
-    image = np.random.randint(low=0, high=255, size=(4, 4, 3), dtype=np.uint8)
-    color = tf.constant([255, 0, 255], tf.uint8)
+@pytest.mark.parametrize("dtype", _DTYPES - {tf.dtypes.float16})
+def test_shear_x(dtype):
+    image = np.random.randint(low=0, high=255, size=(4, 4, 3)).astype(
+        dtype.as_numpy_dtype
+    )
+    color = tf.constant([255, 0, 255], tf.int32)
     level = tf.random.uniform(shape=(), minval=0, maxval=1)
 
     tf_image = tf.constant(image)
@@ -344,11 +348,13 @@ def test_shear_x():
     np.testing.assert_equal(sheared_img.numpy(), expected_img)
 
 
-# TODO: Parameterize on dtypes
 @pytest.mark.usefixtures("maybe_run_functions_eagerly")
-def test_shear_y():
-    image = np.random.randint(low=0, high=255, size=(4, 4, 3), dtype=np.uint8)
-    color = tf.constant([255, 0, 255], tf.dtypes.uint8)
+@pytest.mark.parametrize("dtype", _DTYPES - {tf.dtypes.float16})
+def test_shear_y(dtype):
+    image = np.random.randint(low=0, high=255, size=(4, 4, 3)).astype(
+        dtype.as_numpy_dtype
+    )
+    color = tf.constant([255, 0, 255], tf.int32)
     level = tf.random.uniform(shape=(), minval=0, maxval=1)
 
     tf_image = tf.constant(image)
@@ -363,4 +369,4 @@ def test_shear_y():
     mask = np.where(expected_img == -1)
     expected_img[mask[0], mask[1], :] = color
 
-    np.testing.assert_equal(sheared_img.numpy(), expected_img)
+    test_utils.assert_allclose_according_to_type(sheared_img.numpy(), expected_img)
diff --git a/tensorflow_addons/image/tests/translate_ops_test.py b/tensorflow_addons/image/tests/translate_ops_test.py
--- a/tensorflow_addons/image/tests/translate_ops_test.py
+++ b/tensorflow_addons/image/tests/translate_ops_test.py
@@ -14,13 +14,16 @@
 # ==============================================================================
 """Tests for translate ops."""
 
+import numpy as np
 import pytest
+import scipy
 import tensorflow as tf
-import numpy as np
 
-from tensorflow_addons.image import translate_ops
 from PIL import Image
 
+from tensorflow_addons.image import translate_ops
+from tensorflow_addons.utils import test_utils
+
 _DTYPES = {
     tf.dtypes.uint8,
     tf.dtypes.int32,
@@ -52,7 +55,6 @@ def test_translations_to_projective_transforms():
     np.testing.assert_equal(transform.numpy(), [[1, 0, 1, 0, 1, 1, 0, 0]])
 
 
-# TODO: Parameterize on dtypes
 @pytest.mark.usefixtures("maybe_run_functions_eagerly")
 def test_translate_xy():
     image = np.random.randint(low=0, high=255, size=(4, 4, 3), dtype=np.uint8)
@@ -74,3 +76,23 @@ def test_translate_xy():
     )
 
     np.testing.assert_equal(translated.numpy(), expected)
+
+
+@pytest.mark.usefixtures("maybe_run_functions_eagerly")
+@pytest.mark.parametrize("dtype", _DTYPES - {tf.dtypes.float16})
+def test_translate_xy_scalar_replace(dtype):
+    image = np.random.randint(low=0, high=128, size=(4, 4, 3)).astype(
+        dtype.as_numpy_dtype
+    )
+    translate_to = np.random.randint(low=0, high=4, size=(2,))
+    result = translate_ops.translate_xy(
+        image=image, translate_to=translate_to, replace=1
+    )
+    expected = scipy.ndimage.shift(
+        input=image,
+        shift=(translate_to[1], translate_to[0], 0),
+        order=0,
+        mode="constant",
+        cval=1,
+    )
+    test_utils.assert_allclose_according_to_type(result.numpy(), expected)