[WIP] add autocontrast

tensorflow_addons/image/BUILD

@@ -18,6 +18,7 @@ py_library(
         "connected_components.py",
         "resampler_ops.py",
         "compose_ops.py",
+        "contrast_ops.py",
     ]),
     data = [
         ":sparse_image_warp_test_data",
@@ -177,3 +178,15 @@ py_test(
         ":image",
     ],
 )
+
+py_test(
+    name = "contrast_ops_test",
+    size = "small",
+    srcs = [
+        "contrast_ops_test.py",
+    ],
+    main = "contrast_ops_test.py",
+    deps = [
+        ":image",
+    ],
+)

tensorflow_addons/image/__init__.py

@@ -29,3 +29,4 @@
 from tensorflow_addons.image.transform_ops import transform
 from tensorflow_addons.image.translate_ops import translate
 from tensorflow_addons.image.compose_ops import blend
+from tensorflow_addons.image.contrast_ops import autocontrast

tensorflow_addons/image/contrast_ops.py

@@ -0,0 +1,54 @@
+# 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.
+# ==============================================================================
+import tensorflow as tf
+from tensorflow_addons.utils.types import TensorLike
+
+
+def autocontrast(image: TensorLike) -> tf.Tensor:
+    """Implements Autocontrast function from PIL using TF ops.
+
+    Args:
+      image: A 3D uint8 tensor.
+
+    Returns:
+      The image after it has had autocontrast applied to it and will be of type
+      uint8.
+    """
+    dtype = image.dtype
+    interval = dtype.max - dtype.min
+
+    def scale_channel(image_channel):
+        """Scale the 2D image using the autocontrast rule."""
+        # A possibly cheaper version can be done using cumsum/unique_with_counts
+        # over the histogram values, rather than iterating over the entire image.
+        # to compute mins and maxes.
+        lo = tf.reduce_min(image_channel)
+        hi = tf.reduce_max(image_channel)
+
+        if hi > lo:
+            image_channel = tf.cast(image_channel, tf.float32)
+            image_channel = image_channel * (interval / (hi - lo))
+            image_channel = tf.cast(image_channel, dtype)
+        return image_channel
+
+    # Assumes RGB for now.  Scales each channel independently
+    # and then stacks the result.
+
+    ss = tf.unstack(image, axis=-1)
+    for i in range(len(ss)):
+        ss[i] = scale_channel(ss[i])
+    image = tf.stack(ss, -1)
+
+    return image

tensorflow_addons/image/contrast_ops_test.py

@@ -0,0 +1,43 @@
+# 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.
+# ==============================================================================
+"""Tests for contrast ops."""
+
+import sys
+
+import pytest
+import tensorflow as tf
+import numpy as np
+from absl.testing import parameterized
+from tensorflow_addons.image.contrast_ops import autocontrast
+
+
+@parameterized.named_parameters(
+    ("float16", np.float16), ("float32", np.float32), ("uint8", np.uint8)
+)
+def test_different_dtypes(dtype):
+    test_image = tf.ones([1, 40, 40, 3], dtype=dtype)
+    result_image = autocontrast(test_image)
+    np.testing.assert_allclose(result_image, test_image)
+
+
+def test_different_channels():
+    for channel in [1, 3, 4]:
+        test_image = tf.ones([1, 40, 40, channel], dtype=np.uint8)
+        result_image = autocontrast(test_image)
+        np.testing.assert_allclose(result_image, test_image)
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__]))