Add equalize image op (#1413)

Author: NoobMaster

.github/CODEOWNERS

@@ -36,6 +36,7 @@
 /tensorflow_addons/callbacks/time_stopping*.py @shun-lin
 /tensorflow_addons/callbacks/tqdm_progress_bar*.py @shun-lin
 
+/tensorflow_addons/image/color_ops*.py @abhichou4
 /tensorflow_addons/image/connected_components*.py @sayoojbk
 /tensorflow_addons/image/cutout_ops*.py @fsx950223
 /tensorflow_addons/image/dense_image_warp*.py @windQAQ

tensorflow_addons/image/BUILD

@@ -6,19 +6,20 @@ py_library(
     name = "image",
     srcs = ([
         "__init__.py",
+        "color_ops.py",
+        "compose_ops.py",
+        "connected_components.py",
+        "cutout_ops.py",
         "dense_image_warp.py",
         "distance_transform.py",
         "distort_image_ops.py",
         "filters.py",
+        "interpolate_spline.py",
+        "resampler_ops.py",
+        "sparse_image_warp.py",
         "transform_ops.py",
         "translate_ops.py",
         "utils.py",
-        "sparse_image_warp.py",
-        "interpolate_spline.py",
-        "connected_components.py",
-        "resampler_ops.py",
-        "compose_ops.py",
-        "cutout_ops.py",
     ]),
     data = [
         ":sparse_image_warp_test_data",
@@ -35,157 +36,169 @@ filegroup(
 )
 
 py_test(
-    name = "dense_image_warp_test",
+    name = "color_ops_test",
     size = "small",
     srcs = [
-        "dense_image_warp_test.py",
+        "color_ops_test.py",
     ],
-    main = "dense_image_warp_test.py",
+    main = "color_ops_test.py",
     deps = [
         ":image",
     ],
 )
 
 py_test(
-    name = "distance_transform_ops_test",
-    size = "small",
+    name = "compose_ops_test",
+    size = "medium",
     srcs = [
-        "distance_transform_test.py",
+        "compose_ops_test.py",
     ],
-    main = "distance_transform_test.py",
+    main = "compose_ops_test.py",
     deps = [
         ":image",
     ],
 )
 
 py_test(
-    name = "distort_image_ops_test",
-    size = "small",
+    name = "connected_components_test",
+    size = "medium",
     srcs = [
-        "distort_image_ops_test.py",
+        "connected_components_test.py",
     ],
-    main = "distort_image_ops_test.py",
+    main = "connected_components_test.py",
     deps = [
         ":image",
     ],
 )
 
 py_test(
-    name = "filters_test",
-    size = "medium",
+    name = "cutout_ops_test",
+    size = "small",
     srcs = [
-        "filters_test.py",
+        "cutout_ops_test.py",
     ],
-    flaky = True,
-    main = "filters_test.py",
+    main = "cutout_ops_test.py",
     deps = [
         ":image",
     ],
 )
 
 py_test(
-    name = "transform_ops_test",
-    size = "medium",
+    name = "dense_image_warp_test",
+    size = "small",
     srcs = [
-        "transform_ops_test.py",
+        "dense_image_warp_test.py",
     ],
-    main = "transform_ops_test.py",
+    main = "dense_image_warp_test.py",
     deps = [
         ":image",
     ],
 )
 
 py_test(
-    name = "translate_ops_test",
-    size = "medium",
+    name = "distance_transform_ops_test",
+    size = "small",
     srcs = [
-        "translate_ops_test.py",
+        "distance_transform_test.py",
     ],
-    main = "translate_ops_test.py",
+    main = "distance_transform_test.py",
     deps = [
         ":image",
     ],
 )
 
 py_test(
-    name = "utils_test",
+    name = "distort_image_ops_test",
     size = "small",
     srcs = [
-        "utils_test.py",
+        "distort_image_ops_test.py",
     ],
-    main = "utils_test.py",
+    main = "distort_image_ops_test.py",
     deps = [
         ":image",
     ],
 )
 
 py_test(
-    name = "cutout_ops_test",
-    size = "small",
+    name = "filters_test",
+    size = "medium",
     srcs = [
-        "cutout_ops_test.py",
+        "filters_test.py",
     ],
-    main = "cutout_ops_test.py",
+    flaky = True,
+    main = "filters_test.py",
     deps = [
         ":image",
     ],
 )
 
 py_test(
-    name = "sparse_image_warp_test",
+    name = "interpolate_spline_test",
     size = "medium",
     srcs = [
-        "sparse_image_warp_test.py",
+        "interpolate_spline_test.py",
     ],
-    main = "sparse_image_warp_test.py",
+    main = "interpolate_spline_test.py",
     deps = [
         ":image",
     ],
 )
 
 py_test(
-    name = "interpolate_spline_test",
+    name = "resampler_ops_test",
     size = "medium",
     srcs = [
-        "interpolate_spline_test.py",
+        "resampler_ops_test.py",
     ],
-    main = "interpolate_spline_test.py",
+    main = "resampler_ops_test.py",
     deps = [
         ":image",
     ],
 )
 
 py_test(
-    name = "connected_components_test",
+    name = "sparse_image_warp_test",
     size = "medium",
     srcs = [
-        "connected_components_test.py",
+        "sparse_image_warp_test.py",
     ],
-    main = "connected_components_test.py",
+    main = "sparse_image_warp_test.py",
     deps = [
         ":image",
     ],
 )
 
 py_test(
-    name = "resampler_ops_test",
+    name = "transform_ops_test",
     size = "medium",
     srcs = [
-        "resampler_ops_test.py",
+        "transform_ops_test.py",
     ],
-    main = "resampler_ops_test.py",
+    main = "transform_ops_test.py",
     deps = [
         ":image",
     ],
 )
 
 py_test(
-    name = "compose_ops_test",
+    name = "translate_ops_test",
     size = "medium",
     srcs = [
-        "compose_ops_test.py",
+        "translate_ops_test.py",
     ],
-    main = "compose_ops_test.py",
+    main = "translate_ops_test.py",
+    deps = [
+        ":image",
+    ],
+)
+
+py_test(
+    name = "utils_test",
+    size = "small",
+    srcs = [
+        "utils_test.py",
+    ],
+    main = "utils_test.py",
     deps = [
         ":image",
     ],

tensorflow_addons/image/__init__.py

@@ -14,22 +14,24 @@
 # ==============================================================================
 """Additional image manipulation ops."""
 
+from tensorflow_addons.image.distort_image_ops import adjust_hsv_in_yiq
+from tensorflow_addons.image.compose_ops import blend
+from tensorflow_addons.image.color_ops import equalize
 from tensorflow_addons.image.connected_components import connected_components
+from tensorflow_addons.image.cutout_ops import cutout
 from tensorflow_addons.image.dense_image_warp import dense_image_warp
-from tensorflow_addons.image.dense_image_warp import interpolate_bilinear
 from tensorflow_addons.image.distance_transform import euclidean_dist_transform
-from tensorflow_addons.image.distort_image_ops import adjust_hsv_in_yiq
-from tensorflow_addons.image.distort_image_ops import random_hsv_in_yiq
+from tensorflow_addons.image.dense_image_warp import interpolate_bilinear
+from tensorflow_addons.image.interpolate_spline import interpolate_spline
 from tensorflow_addons.image.filters import mean_filter2d
 from tensorflow_addons.image.filters import median_filter2d
-from tensorflow_addons.image.interpolate_spline import interpolate_spline
+from tensorflow_addons.image.cutout_ops import random_cutout
+from tensorflow_addons.image.distort_image_ops import random_hsv_in_yiq
 from tensorflow_addons.image.resampler_ops import resampler
-from tensorflow_addons.image.sparse_image_warp import sparse_image_warp
 from tensorflow_addons.image.transform_ops import rotate
-from tensorflow_addons.image.transform_ops import transform
 from tensorflow_addons.image.transform_ops import shear_x
 from tensorflow_addons.image.transform_ops import shear_y
+from tensorflow_addons.image.sparse_image_warp import sparse_image_warp
+from tensorflow_addons.image.transform_ops import transform
 from tensorflow_addons.image.translate_ops import translate
 from tensorflow_addons.image.translate_ops import translate_xy
-from tensorflow_addons.image.compose_ops import blend
-from tensorflow_addons.image.cutout_ops import random_cutout, cutout

tensorflow_addons/image/color_ops.py

@@ -0,0 +1,100 @@
+# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Color operations.
+    equalize: Equalizes image histogram
+"""
+
+import tensorflow as tf
+
+from tensorflow_addons.utils.types import TensorLike
+from tensorflow_addons.image.utils import to_4D_image, from_4D_image
+
+from typing import Optional
+from functools import partial
+
+
+def equalize_image(image: TensorLike, data_format: str = "channels_last") -> tf.Tensor:
+    """Implements Equalize function from PIL using TF ops."""
+
+    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: TensorLike, data_format: str = "channels_last", 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) (HWC), or
+          (num_rows, num_columns) (HW). The rank must be statically known (the
+          shape is not `TensorShape(None)`).
+      data_format: Either 'channels_first' or 'channels_last'
+      name: The name of the op.
+    Returns:
+      Image(s) with the same type and shape as `images`, equalized.
+    """
+    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)