Speed up mean_filter2d with depthwise convolution

[WindQAQ]

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 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,

dst: destination array of the same size and type as src.

• For scipy,

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,

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 which supports padding mode. The default padding mode is "REFLECT" because not only scipy but also opencv set REFLECT padding as a default (while some toolboxes in matlab adopt zero-padding by default). Besides, tf.image.sobel_edges, 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
• scipy.ndimage.uniform_filter

[WindQAQ]

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

[Mainak431]

Hmm. As this is faster, implementation can be changed with depthwise conv2d. Returning image to original dtype and range can be cool. But, it will loss the uniformity in output. Most of the functions implemented keeps uniformity of dtype in returned data. You are proposing a single function for single image and batch processing both.What if the first dimension or the batch dimension shape is None or unknown in graph mode?. Did you check for this scenario?

[WindQAQ]

1. It can handle None shape

tf.compat.v1.disable_eager_execution()

x = tf.compat.v1.placeholder(tf.float32, shape=(None, 224, 224, 3))

y = mean_filter2d(x)

with tf.compat.v1.Session() as sess:
    sess.run(y, feed_dict={x: np.random.uniform(size=(64, 224, 224, 3))})


with tf.compat.v1.Session() as sess:
    sess.run(y, feed_dict={x: np.random.uniform(size=(32, 224, 224, 3))})


x2 = tf.compat.v1.placeholder(tf.float32, shape=(224, 224, 3))
y2 = mean_filter2d(x2)

with tf.compat.v1.Session() as sess:
    sess.run(y2, feed_dict={x2: np.random.uniform(size=(224, 224, 3))})

2. Yes, it will loss the uniformity in output, this is only for mean_filter2d though. While this also occurs in scipy, the mean value is computed in float, and recovery it back to the original dtype if needed.
   For median operation, we do not have to cast back to the original dtype since all operations can be done without explicit casting.

[Mainak431]

Ok. Nice. Can you please post the output with real image files. I mean show the converted images. One image pixel value range should be in the range 0 to 255, and another from 0.0 to 1.0. Both kind of images are available in the internet. That will be really helpful. Also posting two converted images in batch processing will ensure the functionality.

[WindQAQ]

Here it is :-)
https://colab.research.google.com/drive/11sz17RCW_L_gJ2kYwCs3f0bR9si3n8PJ
reference here:
https://docs.scipy.org/doc/scipy/reference/generated/scipy.ndimage.uniform_filter.html#scipy.ndimage.uniform_filter

[Mainak431]

Nice. Would you like to implement it? it will also by default solve #230.