Fix unwrap dtype (tensorflow#2121)

Author: WindQAQ

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)

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)

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:

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

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)