Repository: tensorflow/addons

Issue:
Ideally the bins in histogram equalization is variable rather than limited to 256 
https://github.com/tensorflow/addons/blob/d26e2ed5f68092aed57016a7005ce534b1be3dce/tensorflow_addons/image/color_ops.py#L36


Discussion:
It could be but some other APIs don't expose this. E.g. Opencv https://docs.opencv.org/4.5.1/d6/dc7/group__imgproc__hist.html#ga7e54091f0c937d49bf84152a16f76d6e
I can appreciate that the original goal of histeq is to use all the dynamic range for human eyes. The main benefit of adopting 2^12 or 2^16 bins is for where a larger dynamic range worth of signal is available to be standardized for model training purposes. Medical imaging would be where this use case is most commonly used.
You could try to propose a PR so we could see if the new API could be ok.

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,17 +27,17 @@
 from functools import partial
 
 
-def _scale_channel(image: TensorLike, channel: int) -> tf.Tensor:
+def _scale_channel(image: TensorLike, channel: int, bins: int = 256) -> tf.Tensor:
     """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)
+    histo = tf.histogram_fixed_width(image, [0, bins - 1], nbins=bins)
 
     # 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
+    step = (tf.reduce_sum(nonzero_histo) - nonzero_histo[-1]) // (bins - 1)
 
     # If step is zero, return the original image.  Otherwise, build
     # lut from the full histogram and step and then index from it.
@@ -45,20 +45,24 @@ def _scale_channel(image: TensorLike, channel: int) -> tf.Tensor:
         result = image
     else:
         lut_values = (tf.cumsum(histo, exclusive=True) + (step // 2)) // step
-        lut_values = tf.clip_by_value(lut_values, 0, 255)
+        lut_values = tf.clip_by_value(lut_values, 0, bins - 1)
         result = tf.gather(lut_values, image)
 
     return tf.cast(result, image_dtype)
 
 
-def _equalize_image(image: TensorLike) -> tf.Tensor:
+def _equalize_image(image: TensorLike, bins: int = 256) -> 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)
+    image = tf.stack(
+        [_scale_channel(image, c, bins) for c in range(image.shape[-1])], -1
+    )
     return image
 
 
 @tf.function
-def equalize(image: TensorLike, name: Optional[str] = None) -> tf.Tensor:
+def equalize(
+    image: TensorLike, bins: int = 256, name: Optional[str] = None
+) -> tf.Tensor:
     """Equalize image(s)
 
     Args:
@@ -67,6 +71,7 @@ def equalize(image: TensorLike, name: Optional[str] = None) -> tf.Tensor:
           `(num_rows, num_columns, num_channels)` (HWC), or
           `(num_rows, num_columns)` (HW). The rank must be statically known (the
           shape is not `TensorShape(None)`).
+      bins: The number of bins in the histogram.
       name: The name of the op.
     Returns:
       Image(s) with the same type and shape as `images`, equalized.
@@ -75,7 +80,7 @@ def equalize(image: TensorLike, name: Optional[str] = None) -> tf.Tensor:
         image_dims = tf.rank(image)
         image = to_4D_image(image)
         fn = partial(_equalize_image)
-        image = tf.map_fn(fn, image)
+        image = tf.map_fn(lambda x: fn(x, bins), image)
         return from_4D_image(image, image_dims)
 
 


Test patch:
diff --git a/tensorflow_addons/image/tests/color_ops_test.py b/tensorflow_addons/image/tests/color_ops_test.py
--- a/tensorflow_addons/image/tests/color_ops_test.py
+++ b/tensorflow_addons/image/tests/color_ops_test.py
@@ -20,6 +20,7 @@
 
 from tensorflow_addons.image import color_ops
 from PIL import Image, ImageOps, ImageEnhance
+from skimage.exposure import equalize_hist
 
 _DTYPES = {
     np.uint8,
@@ -50,6 +51,27 @@ def test_equalize_with_PIL():
     np.testing.assert_equal(color_ops.equalize(tf.constant(image)).numpy(), equalized)
 
 
+@pytest.mark.usefixtures("maybe_run_functions_eagerly")
+@pytest.mark.parametrize("bins", [256, 65536])
+def test_equalize_with_skimage(bins):
+    np.random.seed(0)
+    image = np.random.randint(low=0, high=256, size=(5, 5, 3, 3), dtype=np.uint8)
+    equalized = tf.cast(color_ops.equalize(tf.constant(image), bins), np.float16)
+    equalized = tf.math.divide(equalized, 255)
+    skiequalized = equalize_hist(image, bins)
+    assert (
+        tf.reduce_mean(
+            tf.math.abs(tf.reshape(tf.math.subtract(equalized, skiequalized), [-1]))
+        )
+        <= 0.015
+    ), "Greater than 0.015 mean absolute difference in arrays"
+    # Test asserts a tolerance of at least 0.015 mean difference between the two
+    # arrays due to the difference in implementation between PIL and skimage.
+    # PIL's and skimage's equalized arrays can have a mean difference ranging
+    # from 0.006-0.036 according showcase done in
+    # https://github.com/tensorflow/addons/pull/2381
+
+
 @pytest.mark.parametrize("dtype", _DTYPES)
 @pytest.mark.parametrize("shape", _SHAPES)
 def test_sharpness_dtype_shape(dtype, shape):