Repository: tensorflow/addons

Issue:
Remove @tf.function in tfa.image.equalize
`tfa.image.equalize` uses an internal `scale_channel` function which triggers excessive retracing when using it with datasets and probably in other cases as well. To illustrate this point the code snippet below benchmarks a naive `tf.data` loop once with the current implementation and once with a copy of the code in `tfa.image.equalize` with `@tf.function` removed.

On my machine I get:
```
[Insert a gazillion retracing warnings here]
With @tf.function: 0.6684637546539307
Without @tf.function: 0.08414626121520996
```

```python
import time
import statistics
from functools import partial

import tensorflow as tf
import tensorflow_addons as tfa
import tensorflow_datasets as tfds

from tensorflow_addons.image.utils import to_4D_image, from_4D_image


def equalize_image(image, data_format: str = "channels_last") -> tf.Tensor:
    """Implements Equalize function from PIL using TF ops."""

    # @tf.function  <-- This triggers retracing.
    def scale_channel(image, channel):
        """Scale the data in the channel to implement equalize."""
        image_dtype = image.dtype

        if data_format == "channels_last":
            image = tf.cast(image[:, :, channel], tf.int32)
        elif data_format == "channels_first":
            image = tf.cast(image[channel], tf.int32)
        else:
            raise ValueError(
                "data_format can either be channels_last or channels_first"
            )
        # Compute the histogram of the image channel.
        histo = tf.histogram_fixed_width(image, [0, 255], nbins=256)

        # For the purposes of computing the step, filter out the nonzeros.
        nonzero = tf.where(tf.not_equal(histo, 0))
        nonzero_histo = tf.reshape(tf.gather(histo, nonzero), [-1])
        step = (tf.reduce_sum(nonzero_histo) - nonzero_histo[-1]) // 255

        def build_lut(histo, step):
            # Compute the cumulative sum, shifting by step // 2
            # and then normalization by step.
            lut = (tf.cumsum(histo) + (step // 2)) // step
            # Shift lut, prepending with 0.
            lut = tf.concat([[0], lut[:-1]], 0)
            # Clip the counts to be in range.  This is done
            # in the C code for image.point.
            return tf.clip_by_value(lut, 0, 255)

        # If step is zero, return the original image.  Otherwise, build
        # lut from the full histogram and step and then index from it.

        if step == 0:
            result = image
        else:
            result = tf.gather(build_lut(histo, step), image)

        return tf.cast(result, image_dtype)

    idx = 2 if data_format == "channels_last" else 0
    image = tf.stack([scale_channel(image, c) for c in range(image.shape[idx])], idx)

    return image


def equalize(
    image, data_format: str = "channels_last", name=None
) -> tf.Tensor:
    with tf.name_scope(name or "equalize"):
        image_dims = tf.rank(image)
        image = to_4D_image(image)
        fn = partial(equalize_image, data_format=data_format)
        image = tf.map_fn(fn, image)
        return from_4D_image(image, image_dims)


dset = tfds.load(name='mnist', data_dir='tensorflow_datasets', split='train')
dset = dset.map(lambda data: data['image']).cache()

# Caching for a fair comparison.
for image in dset:
    img = image
print('Cached dataset')

function_times = list()
for _ in range(10):
    start = time.time()
    for image in dset.take(10):
        augmented_image = tfa.image.equalize(tf.constant(image))
    function_times.append(time.time() - start)
print(f'With @tf.function: {statistics.mean(function_times)}')

non_function_times = list()
for _ in range(10):
    start = time.time()
    for image in dset.take(10):
        augmented_image = equalize(tf.constant(image))
    non_function_times.append(time.time() - start)
print(f'Without @tf.function: {statistics.mean(non_function_times)}')    
```


Discussion:

Resolving patch:
diff --git a/tensorflow_addons/image/color_ops.py b/tensorflow_addons/image/color_ops.py
--- a/tensorflow_addons/image/color_ops.py
+++ b/tensorflow_addons/image/color_ops.py
@@ -27,58 +27,44 @@
 from functools import partial
 
 
-def equalize_image(image: TensorLike) -> tf.Tensor:
-    """Implements Equalize function from PIL using TF ops."""
-
-    @tf.function
-    def scale_channel(image, channel):
-        """Scale the data in the channel to implement equalize."""
-        image_dtype = image.dtype
-        image = tf.cast(image[:, :, channel], tf.int32)
-
-        # Compute the histogram of the image channel.
-        histo = tf.histogram_fixed_width(image, [0, 255], nbins=256)
-
-        # For the purposes of computing the step, filter out the nonzeros.
-        nonzero = tf.where(tf.not_equal(histo, 0))
-        nonzero_histo = tf.reshape(tf.gather(histo, nonzero), [-1])
-        step = (tf.reduce_sum(nonzero_histo) - nonzero_histo[-1]) // 255
+def _scale_channel(image: TensorLike, channel: int) -> tf.Tensor:
+    """Scale the data in the channel to implement equalize."""
+    image_dtype = image.dtype
+    image = tf.cast(image[:, :, channel], tf.int32)
 
-        def build_lut(histo, step):
-            # Compute the cumulative sum, shifting by step // 2
-            # and then normalization by step.
-            lut = (tf.cumsum(histo) + (step // 2)) // step
-            # Shift lut, prepending with 0.
-            lut = tf.concat([[0], lut[:-1]], 0)
-            # Clip the counts to be in range.  This is done
-            # in the C code for image.point.
-            return tf.clip_by_value(lut, 0, 255)
+    # Compute the histogram of the image channel.
+    histo = tf.histogram_fixed_width(image, [0, 255], nbins=256)
 
-        # If step is zero, return the original image.  Otherwise, build
-        # lut from the full histogram and step and then index from it.
+    # For the purposes of computing the step, filter out the nonzeros.
+    nonzero_histo = tf.boolean_mask(histo, histo != 0)
+    step = (tf.reduce_sum(nonzero_histo) - nonzero_histo[-1]) // 255
 
-        if step == 0:
-            result = image
-        else:
-            result = tf.gather(build_lut(histo, step), image)
+    # If step is zero, return the original image.  Otherwise, build
+    # lut from the full histogram and step and then index from it.
+    if step == 0:
+        result = image
+    else:
+        lut_values = (tf.cumsum(histo, exclusive=True) + (step // 2)) // step
+        lut_values = tf.clip_by_value(lut_values, 0, 255)
+        result = tf.gather(lut_values, image)
 
-        return tf.cast(result, image_dtype)
+    return tf.cast(result, image_dtype)
 
-    idx = 2  # Channels last format
-    image = tf.stack([scale_channel(image, c) for c in range(image.shape[idx])], idx)
 
+def _equalize_image(image: TensorLike) -> tf.Tensor:
+    """Implements Equalize function from PIL using TF ops."""
+    image = tf.stack([_scale_channel(image, c) for c in range(image.shape[-1])], -1)
     return image
 
 
+@tf.function
 def equalize(image: TensorLike, name: Optional[str] = None) -> tf.Tensor:
     """Equalize image(s)
 
     Args:
       images: A tensor of shape
           `(num_images, num_rows, num_columns, num_channels)` (NHWC), or
-          `(num_images, num_channels, num_rows, num_columns)` (NCHW), or
           `(num_rows, num_columns, num_channels)` (HWC), or
-          `(num_channels, num_rows, num_columns)` (CHW), or
           `(num_rows, num_columns)` (HW). The rank must be statically known (the
           shape is not `TensorShape(None)`).
       name: The name of the op.
@@ -88,7 +74,7 @@ def equalize(image: TensorLike, name: Optional[str] = None) -> tf.Tensor:
     with tf.name_scope(name or "equalize"):
         image_dims = tf.rank(image)
         image = to_4D_image(image)
-        fn = partial(equalize_image)
+        fn = partial(_equalize_image)
         image = tf.map_fn(fn, image)
         return from_4D_image(image, image_dims)
 


Test patch: