Repository: tensorflow/addons

Issue:
Speed up mean_filter2d with depthwise convolution
**System information**
- TensorFlow version (you are using): **tf-nightly-2.0-preview**
- TensorFlow Addons version: **source**
- Is it in the tf.contrib (if so, where): **no**
- Are you willing to contribute it (yes/no): **yes**
- Are you willing to maintain it going forward? (yes/no): **yes**

**Describe the feature and the current behavior/state.**

Currently, `tfa.image.mean_filter2d` is implemented by `tf.image.extract_patches`, while its functionality is the same with the one applying 2-D convolution with box filter (or uniform filter) channel by channel. This can be easily done by `tf.nn.depthwise_conv2d` with significant speed improvement.

Here is the [notebook example](https://colab.research.google.com/drive/1jJksZAkTfdjJIzPgdw9MhDkNQkFlUM6i) to compare the performance of these two implementations. On colab platform, **it is 4x faster than the original one for either single image or multiple images**. Note that I extend the original implementation with `tf.map_fn` to make it support 4-D input, and some trivial normalization is omitted.

Still wondering why some normalization and casting are performed in the original implementation in both `mean_filter2d` and `median_filter2d`. (no offense, just want to know why)

- For [opencv](https://docs.opencv.org/4.1.0/d4/d86/group__imgproc__filter.html#gad533230ebf2d42509547d514f7d3fbc3),
> dst: destination array of the same size and type as src.
- For [scipy](https://docs.scipy.org/doc/scipy/reference/generated/scipy.ndimage.uniform_filter.html), 
> output : array or dtype, optional
The array in which to place the output, or the dtype of the returned array. By default an array of the same dtype as input will be created.
- For [matlab](https://www.mathworks.com/help/images/ref/medfilt2.html),
> Output image, returned as a numeric matrix of the same class as the input image I.

As far as I'm concerned, there is no need to force the computation limited in range [0, 1], and it's not necessary to transform the output back to `uint8` range even if the input is not of type `uint8`. **All we need to do is to compute the output (may first cast image to float because `depthwise_conv2d` does not accept non-float kernel and input) and cast it back to the original data type.** I believe that for median and average operation, there will not result in any over/underflow situation. Hence, it might be very safe for casting the output back to the original data type, though I think it's better to discuss this first.

Moreover, if users want to do some post processing on the intermediate feature maps from CNNs, these feature maps might not range in either [0, 1] or `uint8`. They are more likely to tensors of real number depending on what the activation function is used.

**Will this change the current api? How?**

`mean_filter2d(image, filter_shape=(3, 3), padding="REFLECT", constant_values=0, name=None)`

The new API is similar to [scipy.ndimage.uniform_filter](https://docs.scipy.org/doc/scipy/reference/generated/scipy.ndimage.uniform_filter.html) which supports padding mode. The default padding mode is "REFLECT" because not only scipy but also [opencv](https://docs.opencv.org/4.1.0/d4/d86/group__imgproc__filter.html#gad533230ebf2d42509547d514f7d3fbc3) set REFLECT padding as a default (while some toolboxes in matlab adopt zero-padding by default). Besides, [tf.image.sobel_edges](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/ops/image_ops_impl.py#L3183), a filter-based application, also uses REFLECT padding mode.

**Who will benefit with this feature?**

People who want to process images.

**Any Other info.**

To conclude:

1. Implement with `tf.nn.depthwise_conv2d`, which can speed up a lot and support batch-wise computation.
2. Cast the input image to float if needed and cast the output back to `image.dtype`.
3. Support padding modes.

Some references for API design:
- [tf.pad](https://www.tensorflow.org/versions/r2.0/api_docs/python/tf/pad)
- [scipy.ndimage.uniform_filter](https://docs.scipy.org/doc/scipy/reference/generated/scipy.ndimage.uniform_filter.html)


Discussion:
Maybe related to #230. Gentle ping to @Mainak431.

Resolving patch:
diff --git a/tensorflow_addons/image/filters.py b/tensorflow_addons/image/filters.py
--- a/tensorflow_addons/image/filters.py
+++ b/tensorflow_addons/image/filters.py
@@ -18,6 +18,7 @@
 from __future__ import print_function
 
 import tensorflow as tf
+from tensorflow_addons.utils import keras_utils
 
 
 @tf.function
@@ -34,82 +35,110 @@ def func2():
     return tf.cond(tf.math.greater(ma, one), func2, func1)
 
 
-@tf.function
-def mean_filter2d(image, filter_shape=(3, 3), name=None):
-    """This method performs Mean Filtering on image. Filter shape can be user
-    given.
+def _pad(image, filter_shape, mode="CONSTANT", constant_values=0):
+    """Explicitly pad a 4-D image.
+
+    Equivalent to the implicit padding method offered in `tf.nn.conv2d` and
+    `tf.nn.depthwise_conv2d`, but supports non-zero, reflect and symmetric
+    padding mode. For the even-sized filter, it pads one more value to the
+    right or the bottom side.
 
-    This method takes both kind of images where pixel values lie between 0 to
-    255 and where it lies between 0.0 and 1.0
     Args:
-        image: A 3D `Tensor` of type `float32` or 'int32' or 'float64' or
-               'int64 and of shape`[rows, columns, channels]`
+      image: A 4-D `Tensor` of shape `[batch_size, height, width, channels]`.
+      filter_shape: A `tuple`/`list` of 2 integers, specifying the height
+        and width of the 2-D filter.
+      mode: A `string`, one of "REFLECT", "CONSTANT", or "SYMMETRIC".
+        The type of padding algorithm to use, which is compatible with
+        `mode` argument in `tf.pad`. For more details, please refer to
+        https://www.tensorflow.org/api_docs/python/tf/pad.
+      constant_values: A `scalar`, the pad value to use in "CONSTANT"
+        padding mode.
+    """
+    assert mode in ["CONSTANT", "REFLECT", "SYMMETRIC"]
+    filter_height, filter_width = filter_shape
+    pad_top = (filter_height - 1) // 2
+    pad_bottom = filter_height - 1 - pad_top
+    pad_left = (filter_width - 1) // 2
+    pad_right = filter_width - 1 - pad_left
+    paddings = [[0, 0], [pad_top, pad_bottom], [pad_left, pad_right], [0, 0]]
+    return tf.pad(image, paddings, mode=mode, constant_values=constant_values)
 
-        filter_shape: Optional Argument. A tuple of 2 integers (R,C).
-               R is the first value is the number of rows in the filter and
-               C is the second value in the filter is the number of columns
-               in the filter. This creates a filter of shape (R,C) or RxC
-               filter. Default value = (3,3)
 
-     Returns:
-         A 3D mean filtered image tensor of shape [rows,columns,channels] and
-         type 'int32'. Pixel value of returned tensor ranges between 0 to 255
-    """
+@tf.function
+def mean_filter2d(image,
+                  filter_shape=(3, 3),
+                  padding="REFLECT",
+                  constant_values=0,
+                  name=None):
+    """Perform mean filtering on image(s).
 
+    Args:
+      image: Either 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
+        to specify the same value for all spatial dimensions.
+      padding: A `string`, one of "REFLECT", "CONSTANT", or "SYMMETRIC".
+        The type of padding algorithm to use, which is compatible with
+        `mode` argument in `tf.pad`. For more details, please refer to
+        https://www.tensorflow.org/api_docs/python/tf/pad.
+      constant_values: A `scalar`, the pad value to use in "CONSTANT"
+        padding mode.
+      name: A name for this operation (optional).
+    Returns:
+      3-D or 4-D `Tensor` of the same dtype as input.
+    Raises:
+      ValueError: If `image` is not 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"):
-        if not isinstance(filter_shape, tuple):
-            raise TypeError('Filter shape must be a tuple')
-        if len(filter_shape) != 2:
-            raise ValueError('Filter shape must be a tuple of 2 integers. '
-                             'Got %s values in tuple' % len(filter_shape))
-        filter_shapex = filter_shape[0]
-        filter_shapey = filter_shape[1]
-        if not isinstance(filter_shapex, int) or not isinstance(
-                filter_shapey, int):
-            raise TypeError('Size of the filter must be Integers')
-        (row, col, ch) = (image.shape[0], image.shape[1], image.shape[2])
-        if row != None and col != None and ch != None:
-            (row, col, ch) = (int(row), int(col), int(ch))
-        else:
-            raise TypeError(
-                'All the Dimensions of the input image tensor must be \
-                Integers.')
-        if row < filter_shapex or col < filter_shapey:
+        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.")
+
+        if padding not in ["REFLECT", "CONSTANT", "SYMMETRIC"]:
             raise ValueError(
-                'Number of Pixels in each dimension of the image should be \
-                more than the filter size. Got filter_shape (%sx' %
-                filter_shape[0] + '%s).' % filter_shape[1] +
-                ' Image Shape (%s)' % image.shape)
-        if filter_shapex % 2 == 0 or filter_shapey % 2 == 0:
-            raise ValueError('Filter size should be odd. Got filter_shape (%sx'
-                             % filter_shape[0] + '%s)' % filter_shape[1])
-        image = tf.cast(image, tf.float32)
-        tf_i = tf.reshape(image, [row * col * ch])
-        ma = tf.math.reduce_max(tf_i)
-        image = _normalize(image, ma)
+                "padding should be one of \"REFLECT\", \"CONSTANT\", or "
+                "\"SYMMETRIC\".")
 
-        # k and l is the Zero-padding size
+        filter_shape = keras_utils.conv_utils.normalize_tuple(
+            filter_shape, 2, "filter_shape")
 
-        listi = []
-        for a in range(ch):
-            img = image[:, :, a:a + 1]
-            img = tf.reshape(img, [1, row, col, 1])
-            slic = tf.image.extract_patches(
-                img, [1, filter_shapex, filter_shapey, 1], [1, 1, 1, 1],
-                [1, 1, 1, 1],
-                padding='SAME')
-            li = tf.reduce_mean(slic, axis=-1)
-            li = tf.reshape(li, [row, col, 1])
-            listi.append(li)
-        y = tf.concat(listi[0], 2)
+        # Expand to a 4-D tensor
+        if rank == 3:
+            image = tf.expand_dims(image, axis=0)
 
-        for i in range(len(listi) - 1):
-            y = tf.concat([y, listi[i + 1]], 2)
+        # Keep the precision if it's float;
+        # otherwise, convert to float32 for computing.
+        orig_dtype = image.dtype
+        if not image.dtype.is_floating:
+            image = tf.dtypes.cast(image, tf.dtypes.float32)
 
-        y *= 255
-        y = tf.cast(y, tf.int32)
+        # Explicitly pad the image
+        image = _pad(
+            image, filter_shape, mode=padding, constant_values=constant_values)
 
-        return y
+        # Filter of shape (filter_width, filter_height, in_channels, 1)
+        # has the value of 1 for each element.
+        area = tf.constant(
+            filter_shape[0] * filter_shape[1], dtype=image.dtype)
+        filter_shape = filter_shape + (tf.shape(image)[-1], 1)
+        kernel = tf.ones(shape=filter_shape, dtype=image.dtype)
+
+        output = tf.nn.depthwise_conv2d(
+            image, kernel, strides=(1, 1, 1, 1), padding="VALID")
+
+        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)
+
+        return tf.dtypes.cast(output, orig_dtype)
 
 
 @tf.function
diff --git a/tensorflow_addons/utils/keras_utils.py b/tensorflow_addons/utils/keras_utils.py
--- a/tensorflow_addons/utils/keras_utils.py
+++ b/tensorflow_addons/utils/keras_utils.py
@@ -21,6 +21,7 @@
 
 # TODO: find public API alternative to these
 from tensorflow.python.keras.losses import LossFunctionWrapper  # pylint: disable=unused-import
+from tensorflow.python.keras.utils import conv_utils  # pylint: disable=unused-import
 
 
 def register_keras_custom_object(cls):


Test patch:
diff --git a/tensorflow_addons/image/filters_test.py b/tensorflow_addons/image/filters_test.py
--- a/tensorflow_addons/image/filters_test.py
+++ b/tensorflow_addons/image/filters_test.py
@@ -23,73 +23,159 @@
 from tensorflow_addons.utils import test_utils
 
 
+@test_utils.run_all_in_graph_and_eager_modes
 class MeanFilter2dTest(tf.test.TestCase):
-    def _validate_mean_filter2d(self,
-                                inputs,
-                                expected_values,
-                                filter_shape=(3, 3)):
-        output = mean_filter2d(inputs, filter_shape)
-        self.assertAllClose(output, expected_values)
+    def _tile_image(self, plane, image_shape):
+        """Tile a 2-D image `plane` into 3-D or 4-D as per `image_shape`."""
+        assert 3 <= len(image_shape) <= 4
+        plane = tf.convert_to_tensor(plane)
+        plane = tf.expand_dims(plane, -1)
+        channels = image_shape[-1]
+        image = tf.tile(plane, (1, 1, channels))
+
+        if len(image_shape) == 4:
+            batch_size = image_shape[0]
+            image = tf.expand_dims(image, 0)
+            image = tf.tile(image, (batch_size, 1, 1, 1))
+
+        return image
+
+    def _setup_values(self, image_shape, filter_shape, padding,
+                      constant_values, dtype):
+        assert 3 <= len(image_shape) <= 4
+        height, width = image_shape[-3], image_shape[-2]
+        plane = tf.constant([x for x in range(1, height * width + 1)],
+                            shape=(height, width),
+                            dtype=dtype)
+        image = self._tile_image(plane, image_shape=image_shape)
+
+        result = mean_filter2d(
+            image,
+            filter_shape=filter_shape,
+            padding=padding,
+            constant_values=constant_values)
+
+        return result
+
+    def _verify_values(self, image_shape, filter_shape, padding,
+                       constant_values, expected_plane):
+
+        expected_output = self._tile_image(expected_plane, image_shape)
+        dtypes = tf.dtypes
+        for dtype in [
+                dtypes.uint8, dtypes.float16, dtypes.float32, dtypes.float64
+        ]:
+            result = self._setup_values(image_shape, filter_shape, padding,
+                                        constant_values, dtype)
+            self.assertAllCloseAccordingToType(
+                result, tf.dtypes.cast(expected_output, dtype))
+
+    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)]:
+            with self.subTest(dim=len(image_shape)):
+                with self.assertRaisesRegexp(ValueError, msg):
+                    mean_filter2d(tf.ones(shape=image_shape))
+
+    def test_invalid_filter_shape(self):
+        msg = ("The `filter_shape` argument must be a tuple of 2 integers.")
+        image = tf.ones(shape=(1, 28, 28, 1))
+
+        for filter_shape in [(3, 3, 3), (3, None, 3), None]:
+            with self.subTest(filter_shape=filter_shape):
+                with self.assertRaisesRegexp(ValueError, msg):
+                    mean_filter2d(image, filter_shape=filter_shape)
+
+    def test_invalid_padding(self):
+        msg = ("padding should be one of \"REFLECT\", \"CONSTANT\", "
+               "or \"SYMMETRIC\".")
+        image = tf.ones(shape=(1, 28, 28, 1))
 
-    @test_utils.run_in_graph_and_eager_modes
-    def test_filter_tuple(self):
-        tf_img = tf.zeros([3, 4, 3], tf.int32)
-
-        for filter_shape in [3, 3.5, 'dt', None]:
-            with self.assertRaisesRegexp(TypeError,
-                                         'Filter shape must be a tuple'):
-                mean_filter2d(tf_img, filter_shape)
-
-        filter_shape = (3, 3, 3)
-        msg = ('Filter shape must be a tuple of 2 integers. '
-               'Got %s values in tuple' % len(filter_shape))
         with self.assertRaisesRegexp(ValueError, msg):
-            mean_filter2d(tf_img, filter_shape)
-
-        msg = 'Size of the filter must be Integers'
-        for filter_shape in [(3.5, 3), (None, 3)]:
-            with self.assertRaisesRegexp(TypeError, msg):
-                mean_filter2d(tf_img, filter_shape)
-
-    @test_utils.run_in_graph_and_eager_modes
-    def test_filter_value(self):
-        tf_img = tf.zeros([3, 4, 3], tf.int32)
-
-        with self.assertRaises(ValueError):
-            mean_filter2d(tf_img, (4, 3))
-
-    @test_utils.run_deprecated_v1
-    def test_dimension(self):
-        for image_shape in [(3, 4, None), (3, None, 4), (None, 3, 4)]:
-            with self.assertRaises(TypeError):
-                tf_img = tf.compat.v1.placeholder(tf.int32, shape=image_shape)
-                mean_filter2d(tf_img)
-
-    @test_utils.run_in_graph_and_eager_modes
-    def test_image_vs_filter(self):
-        tf_img = tf.zeros([3, 4, 3], tf.int32)
-        filter_shape = (3, 5)
-        with self.assertRaises(ValueError):
-            mean_filter2d(tf_img, filter_shape)
-
-    @test_utils.run_in_graph_and_eager_modes
-    def test_three_channels(self):
-        tf_img = [[[0.32801723, 0.08863795, 0.79119259],
-                   [0.35526001, 0.79388736, 0.55435993],
-                   [0.11607035, 0.55673079, 0.99473371]],
-                  [[0.53240645, 0.74684819, 0.33700031],
-                   [0.01760473, 0.28181609, 0.9751476],
-                   [0.01605137, 0.8292904, 0.56405609]],
-                  [[0.57215374, 0.10155051, 0.64836128],
-                   [0.36533048, 0.91401874, 0.02524159],
-                   [0.56379134, 0.9028874, 0.19505117]]]
-
-        tf_img = tf.convert_to_tensor(value=tf_img)
-        expt = [[[34, 54, 75], [38, 93, 119], [14, 69, 87]],
-                [[61, 82, 94], [81, 147, 144], [40, 121, 93]],
-                [[42, 57, 56], [58, 106, 77], [27, 82, 49]]]
-        expt = tf.convert_to_tensor(value=expt)
-        self._validate_mean_filter2d(tf_img, expt)
+            mean_filter2d(image, padding="TEST")
+
+    def test_none_channels(self):
+        # 3-D image
+        fn = mean_filter2d.get_concrete_function(
+            tf.TensorSpec(dtype=tf.dtypes.float32, shape=(3, 3, None)))
+        fn(tf.ones(shape=(3, 3, 1)))
+        fn(tf.ones(shape=(3, 3, 3)))
+
+        # 4-D image
+        fn = mean_filter2d.get_concrete_function(
+            tf.TensorSpec(dtype=tf.dtypes.float32, shape=(1, 3, 3, None)))
+        fn(tf.ones(shape=(1, 3, 3, 1)))
+        fn(tf.ones(shape=(1, 3, 3, 3)))
+
+    def test_reflect_padding_with_3x3_filter(self):
+        expected_plane = tf.constant([[33. / 9., 36. / 9., 39. / 9.],
+                                      [42. / 9., 45. / 9., 48. / 9.],
+                                      [51. / 9., 54. / 9., 57. / 9.]])
+
+        for image_shape in [(3, 3, 1), (3, 3, 3), (1, 3, 3, 1), (1, 3, 3, 3),
+                            (2, 3, 3, 1), (2, 3, 3, 3)]:
+            self._verify_values(
+                image_shape=image_shape,
+                filter_shape=(3, 3),
+                padding="REFLECT",
+                constant_values=0,
+                expected_plane=expected_plane)
+
+    def test_reflect_padding_with_4x4_filter(self):
+        expected_plane = tf.constant([[80. / 16., 80. / 16., 80. / 16.],
+                                      [80. / 16., 80. / 16., 80. / 16.],
+                                      [80. / 16., 80. / 16., 80. / 16.]])
+
+        for image_shape in [(3, 3, 1), (3, 3, 3), (1, 3, 3, 1), (1, 3, 3, 3),
+                            (2, 3, 3, 1), (2, 3, 3, 3)]:
+            self._verify_values(
+                image_shape=image_shape,
+                filter_shape=(4, 4),
+                padding="REFLECT",
+                constant_values=0,
+                expected_plane=expected_plane)
+
+    def test_constant_padding_with_3x3_filter(self):
+        expected_plane = tf.constant([[12. / 9., 21. / 9., 16. / 9.],
+                                      [27. / 9., 45. / 9., 33. / 9.],
+                                      [24. / 9., 39. / 9., 28. / 9.]])
+
+        for image_shape in [(3, 3, 1), (3, 3, 3), (1, 3, 3, 1), (1, 3, 3, 3),
+                            (2, 3, 3, 1), (2, 3, 3, 3)]:
+            self._verify_values(
+                image_shape=image_shape,
+                filter_shape=(3, 3),
+                padding="CONSTANT",
+                constant_values=0,
+                expected_plane=expected_plane)
+
+        expected_plane = tf.constant([[17. / 9., 24. / 9., 21. / 9.],
+                                      [30. / 9., 45. / 9., 36. / 9.],
+                                      [29. / 9., 42. / 9., 33. / 9.]])
+
+        for image_shape in [(3, 3, 1), (3, 3, 3), (1, 3, 3, 1), (1, 3, 3, 3),
+                            (2, 3, 3, 1), (2, 3, 3, 3)]:
+            self._verify_values(
+                image_shape=image_shape,
+                filter_shape=(3, 3),
+                padding="CONSTANT",
+                constant_values=1,
+                expected_plane=expected_plane)
+
+    def test_symmetric_padding_with_3x3_filter(self):
+        expected_plane = tf.constant([[21. / 9., 27. / 9., 33. / 9.],
+                                      [39. / 9., 45. / 9., 51. / 9.],
+                                      [57. / 9., 63. / 9., 69. / 9.]])
+
+        for image_shape in [(3, 3, 1), (3, 3, 3), (1, 3, 3, 1), (1, 3, 3, 3),
+                            (2, 3, 3, 1), (2, 3, 3, 3)]:
+            self._verify_values(
+                image_shape=image_shape,
+                filter_shape=(3, 3),
+                padding="SYMMETRIC",
+                constant_values=0,
+                expected_plane=expected_plane)
 
 
 class MedianFilter2dTest(tf.test.TestCase):