[image/filters] convert unknown rank to 4d (#343)

* convert unknown rank to 4d

Author: WindQAQ

tensorflow_addons/image/filters.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 import keras_utils
 
 
@@ -59,7 +60,8 @@ def mean_filter2d(image,
     """Perform mean filtering on image(s).
 
     Args:
-      image: Either a 3-D `Tensor` of shape `[height, width, channels]`,
+      image: Either a 2-D `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]`.
       filter_shape: An `integer` or `tuple`/`list` of 2 integers, specifying
         the height and width of the 2-D mean filter. Can be a single integer
@@ -74,16 +76,14 @@ def mean_filter2d(image,
     Returns:
       3-D or 4-D `Tensor` of the same dtype as input.
     Raises:
-      ValueError: If `image` is not 3 or 4-dimensional,
+      ValueError: If `image` is not 2, 3 or 4-dimensional,
         if `padding` is other than "REFLECT", "CONSTANT" or "SYMMETRIC",
         or if `filter_shape` is invalid.
     """
     with tf.name_scope(name or "mean_filter2d"):
         image = tf.convert_to_tensor(image, name="image")
-
-        rank = image.shape.rank
-        if rank != 3 and rank != 4:
-            raise ValueError("image should be either 3 or 4-dimensional.")
+        original_ndims = img_utils.get_ndims(image)
+        image = img_utils.to_4D_image(image)
 
         if padding not in ["REFLECT", "CONSTANT", "SYMMETRIC"]:
             raise ValueError(
@@ -93,10 +93,6 @@ def mean_filter2d(image,
         filter_shape = keras_utils.normalize_tuple(filter_shape, 2,
                                                    "filter_shape")
 
-        # Expand to a 4-D tensor
-        if rank == 3:
-            image = tf.expand_dims(image, axis=0)
-
         # Keep the precision if it's float;
         # otherwise, convert to float32 for computing.
         orig_dtype = image.dtype
@@ -119,11 +115,7 @@ def mean_filter2d(image,
 
         output /= area
 
-        # Squeeze out the first axis to make sure
-        # output has the same dimension with image.
-        if rank == 3:
-            output = tf.squeeze(output, axis=0)
-
+        output = img_utils.from_4D_image(output, original_ndims)
         return tf.dtypes.cast(output, orig_dtype)
 
 
@@ -136,7 +128,8 @@ def median_filter2d(image,
     """Perform median filtering on image(s).
 
     Args:
-      image: Either a 3-D `Tensor` of shape `[height, width, channels]`,
+      image: Either a 2-D `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]`.
       filter_shape: An `integer` or `tuple`/`list` of 2 integers, specifying
         the height and width of the 2-D median filter. Can be a single integer
@@ -151,16 +144,14 @@ def median_filter2d(image,
     Returns:
       3-D or 4-D `Tensor` of the same dtype as input.
     Raises:
-      ValueError: If `image` is not 3 or 4-dimensional,
+      ValueError: If `image` is not 2, 3 or 4-dimensional,
         if `padding` is other than "REFLECT", "CONSTANT" or "SYMMETRIC",
         or if `filter_shape` is invalid.
     """
     with tf.name_scope(name or "median_filter2d"):
         image = tf.convert_to_tensor(image, name="image")
-
-        rank = image.shape.rank
-        if rank != 3 and rank != 4:
-            raise ValueError("image should be either 3 or 4-dimensional.")
+        original_ndims = img_utils.get_ndims(image)
+        image = img_utils.to_4D_image(image)
 
         if padding not in ["REFLECT", "CONSTANT", "SYMMETRIC"]:
             raise ValueError(
@@ -170,10 +161,6 @@ def median_filter2d(image,
         filter_shape = keras_utils.normalize_tuple(filter_shape, 2,
                                                    "filter_shape")
 
-        # Expand to a 4-D tensor
-        if rank == 3:
-            image = tf.expand_dims(image, axis=0)
-
         image_shape = tf.shape(image)
         batch_size = image_shape[0]
         height = image_shape[1]
@@ -212,10 +199,5 @@ def median_filter2d(image,
                 top[:, :, :, :, floor - 1] + top[:, :, :, :, ceil - 1]) / 2
 
         output = tf.cast(median, image.dtype)
-
-        # Squeeze out the first axis to make sure
-        # output has the same dimension with image.
-        if rank == 3:
-            output = tf.squeeze(output, axis=0)
-
+        output = img_utils.from_4D_image(output, original_ndims)
         return output

tensorflow_addons/image/filters_test.py

@@ -82,12 +82,13 @@ def setUp(self):
         super(MeanFilter2dTest, self).setUp()
 
     def test_invalid_image(self):
-        msg = "image should be either 3 or 4-dimensional."
-
-        for image_shape in [(28, 28), (16, 28, 28, 1, 1)]:
+        msg = "`image` must be 2/3/4D tensor"
+        errors = (ValueError, tf.errors.InvalidArgumentError)
+        for image_shape in [(1,), (16, 28, 28, 1, 1)]:
             with self.subTest(dim=len(image_shape)):
-                with self.assertRaisesRegexp(ValueError, msg):
-                    mean_filter2d(tf.ones(shape=image_shape))
+                with self.assertRaisesRegexp(errors, msg):
+                    image = tf.ones(shape=image_shape)
+                    self.evaluate(mean_filter2d(image))
 
     def test_invalid_filter_shape(self):
         msg = ("The `filter_shape` argument must be a tuple of 2 integers.")
@@ -119,6 +120,16 @@ def test_none_channels(self):
         fn(tf.ones(shape=(1, 3, 3, 1)))
         fn(tf.ones(shape=(1, 3, 3, 3)))
 
+    def test_unknown_shape(self):
+        fn = mean_filter2d.get_concrete_function(
+            tf.TensorSpec(shape=None, dtype=tf.dtypes.float32),
+            padding="CONSTANT",
+            constant_values=1.)
+
+        for shape in [(3, 3), (3, 3, 3), (1, 3, 3, 3)]:
+            image = tf.ones(shape=shape)
+            self.assertAllEqual(self.evaluate(image), self.evaluate(fn(image)))
+
     def test_reflect_padding_with_3x3_filter(self):
         expected_plane = tf.constant([[33. / 9., 36. / 9., 39. / 9.],
                                       [42. / 9., 45. / 9., 48. / 9.],
@@ -191,12 +202,13 @@ def setUp(self):
         super(MedianFilter2dTest, self).setUp()
 
     def test_invalid_image(self):
-        msg = "image should be either 3 or 4-dimensional."
-
-        for image_shape in [(28, 28), (16, 28, 28, 1, 1)]:
+        msg = "`image` must be 2/3/4D tensor"
+        errors = (ValueError, tf.errors.InvalidArgumentError)
+        for image_shape in [(1,), (16, 28, 28, 1, 1)]:
             with self.subTest(dim=len(image_shape)):
-                with self.assertRaisesRegexp(ValueError, msg):
-                    median_filter2d(tf.ones(shape=image_shape))
+                with self.assertRaisesRegexp(errors, msg):
+                    image = tf.ones(shape=image_shape)
+                    self.evaluate(median_filter2d(image))
 
     def test_invalid_filter_shape(self):
         msg = ("The `filter_shape` argument must be a tuple of 2 integers.")
@@ -228,6 +240,16 @@ def test_none_channels(self):
         fn(tf.ones(shape=(1, 3, 3, 1)))
         fn(tf.ones(shape=(1, 3, 3, 3)))
 
+    def test_unknown_shape(self):
+        fn = median_filter2d.get_concrete_function(
+            tf.TensorSpec(shape=None, dtype=tf.dtypes.float32),
+            padding="CONSTANT",
+            constant_values=1.)
+
+        for shape in [(3, 3), (3, 3, 3), (1, 3, 3, 3)]:
+            image = tf.ones(shape=shape)
+            self.assertAllEqual(self.evaluate(image), self.evaluate(fn(image)))
+
     def test_reflect_padding_with_3x3_filter(self):
         expected_plane = tf.constant([[4, 4, 5], [5, 5, 5], [5, 6, 6]])