Repository: tensorflow/addons

Issue:
Inconsistent APIs within tfa.image and between tf.image and tfa.image
**Describe the feature and the current behavior/state.**
Currently, some functions in `tfa.image` that support `data_format` args like [equalize](https://github.com/tensorflow/addons/blob/master/tensorflow_addons/image/color_ops.py#L79), [cutout](https://github.com/tensorflow/addons/blob/master/tensorflow_addons/image/cutout_ops.py#L97), and [random_cutout](https://github.com/tensorflow/addons/blob/master/tensorflow_addons/image/cutout_ops.py#L46), while most of them do not.  Moreover, all functions in [tf.image](https://www.tensorflow.org/api_docs/python/tf/image) support only channel-last images.

Therefore, I wonder if we can drop the support for `data_format` argument and channel-first image in `tfa.image` for some reasons:

- Be consistent with `tf.image`: Usually, we compose multiple functions in `tf.image` and `tfa.image`. So if one of the function does not support channel-first image, there must exist at least one `tf.transpose` operation to change the data format. In this case, the support of channel-first image seems vague and useless as we have to explicitly change it into channel-last format in order for input pipeline.
- Unlike Conv2D or CorrelationCost that has specific custom ops to handle different format, the support for channel first images in `tfa.image` just selects different axis to unpack, compute and pack.
- Shape casting: we do `to_4D_image` and `from_4D_image` to make most operations compatible with `tf.data.Dataset`, while they operate only on channel last images. We can indeed extend those functions for channel first images, but I am also confused if it's worthwhile.

**Relevant information**
- Are you willing to contribute it (yes/no): yes
- Are you willing to maintain it going forward? (yes/no): yes
- Is there a relevant academic paper? (if so, where): no
- Is there already an implementation in another framework? (if so, where): addons
- Was it part of tf.contrib? (if so, where): no

**Which API type would this fall under (layer, metric, optimizer, etc.)**
image

**Who will benefit with this feature?**
developers who use tfa.image

**Any other info.**


Discussion:
I prefer to be uniform with `tf.image` and also with `keras_image` in the future if it will follow this same "standard" format.
> Be consistent with tf.image: Usually, we compose multiple functions in tf.image and tfa.image. So if one of the function does not support channel-first image, there must exist at least one tf.transpose operation to change the data format. In this case, the support of channel-first image seems vague and useless as we have to explicitly change it into channel-last format in order for input pipeline.

I am not sure if we are uniform also about the casting operations.
Thanks for bringing this up!  This much very much needs to be decided upon and standardized. 

I also agree we should be consistant with what `tf.image` supports, but just want to confirm that NHWC is all that tf.image plans to support in the future. Seeing as keras has `data_format` args, I want to verify that support isn't on the roadmap for tf.image ops given the new Keras preprocessing pipelines.

@tanzhenyu Would you be able to confirm that our understanding is correct or point us to someone who could? 
For conv ops:
As of today, NHWC is what CPU only supports. In Cuda, it will mostly convert to NCHW first and call an efficient cudnn, except if it's float16 and > Volta V100, but this is done by default.

For tf.image ops:
They're all based on NHWC format, in the kernel level. This is not something we can change at python level at all.

So yeah I think we should be consistent with tf.image
+1, it sounds great to keep `tfa.image` simple enough.
Thanks for the comments! I also recommend that `tfa.image` should be compatible with 3-D image because the common use case is to use `.map(lambda x: tfa.image.foo(x))` to process a batch of images.
> Shape casting: we do to_4D_image and from_4D_image to make most operations compatible with tf.data.Dataset, while they operate only on channel last images. We can indeed extend those functions for channel first images, but I am also confused if it's worthwhile.

@WindQAQ I don't understand in the current status if we need to do this or not cause in https://github.com/tensorflow/addons/pull/1840 we was going to set the shape bypassing the `4d` utils.
> > Shape casting: we do to_4D_image and from_4D_image to make most operations compatible with tf.data.Dataset, while they operate only on channel last images. We can indeed extend those functions for channel first images, but I am also confused if it's worthwhile.
> 
> @WindQAQ I don't understand in the current status if we need to do this or not cause in #1840 we was going to set the shape bypassing the `4d` utils.

@bhack I suppose we should at least support 3D image as most of functions in `tf.image` do that. So we can have another PR to fix it.
@WindQAQ What Is supporting the official Vision Model garden refactoring  https://github.com/tensorflow/models/blob/master/official/vision/image_classification/augment.py#L37?
Do you think that users could mix that with `tfa.image` or `tf.image`?
> @WindQAQ What Is supporting the official Vision Model garden refactoring https://github.com/tensorflow/models/blob/master/official/vision/image_classification/augment.py#L37?
> Do you think that users could mix that with `tfa.image` or `tf.image`?

When migrating everything to addons (or others repo), it should stick to the rule of addons IMO. So the implementation should be either refactored or discarded. 
> Thanks for the comments! I also recommend that `tfa.image` should be compatible with 3-D image because the common use case is to use `.map(lambda x: tfa.image.foo(x))` to process a batch of images.

In most cases, it might be better to use `data.batch(xx).map(xxx)` for performance reason :-)
You can check Gsoc progress at https://github.com/tanzhenyu/image_augmentation

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
@@ -24,10 +24,9 @@
 from tensorflow_addons.image.compose_ops import blend
 
 from typing import Optional
-from functools import partial
 
 
-def equalize_image(image: TensorLike, data_format: str = "channels_last") -> tf.Tensor:
+def equalize_image(image: TensorLike) -> tf.Tensor:
     """Implements Equalize function from PIL using TF ops."""
 
     @tf.function
@@ -35,14 +34,7 @@ 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"
-            )
+        image = tf.cast(image[:, :, channel], tf.int32)
         # Compute the histogram of the image channel.
         histo = tf.histogram_fixed_width(image, [0, 255], nbins=256)
 
@@ -71,26 +63,20 @@ def build_lut(histo, step):
 
         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)
+    image = tf.stack([scale_channel(image, c) for c in range(image.shape[2])], 2)
 
     return image
 
 
-def equalize(
-    image: TensorLike, data_format: str = "channels_last", name: Optional[str] = None
-) -> tf.Tensor:
+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)`).
-      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.
@@ -98,8 +84,7 @@ def equalize(
     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)
+        image = tf.map_fn(equalize_image, image)
         return from_4D_image(image, image_dims)
 
 
diff --git a/tensorflow_addons/image/cutout_ops.py b/tensorflow_addons/image/cutout_ops.py
--- a/tensorflow_addons/image/cutout_ops.py
+++ b/tensorflow_addons/image/cutout_ops.py
@@ -15,20 +15,11 @@
 """Cutout op"""
 
 import tensorflow as tf
-from tensorflow_addons.utils import keras_utils
-from tensorflow_addons.utils.types import TensorLike, Number
-
 
-def _get_image_wh(images, data_format):
-    if data_format == "channels_last":
-        image_height, image_width = tf.shape(images)[1], tf.shape(images)[2]
-    else:
-        image_height, image_width = tf.shape(images)[2], tf.shape(images)[3]
-
-    return image_height, image_width
+from tensorflow_addons.utils.types import TensorLike, Number
 
 
-def _norm_params(images, mask_size, data_format):
+def _norm_params(images, mask_size):
     mask_size = tf.convert_to_tensor(mask_size)
     if tf.executing_eagerly():
         tf.assert_equal(
@@ -38,9 +29,9 @@ def _norm_params(images, mask_size, data_format):
         )
     if tf.rank(mask_size) == 0:
         mask_size = tf.stack([mask_size, mask_size])
-    data_format = keras_utils.normalize_data_format(data_format)
-    image_height, image_width = _get_image_wh(images, data_format)
-    return mask_size, data_format, image_height, image_width
+    shape = tf.shape(images)
+    image_height, image_width = shape[1], shape[2]
+    return mask_size, image_height, image_width
 
 
 def random_cutout(
@@ -48,7 +39,6 @@ def random_cutout(
     mask_size: TensorLike,
     constant_values: Number = 0,
     seed: Number = None,
-    data_format: str = "channels_last",
 ) -> tf.Tensor:
     """Apply cutout (https://arxiv.org/abs/1708.04552) to images.
 
@@ -58,9 +48,7 @@ def random_cutout(
     chosen uniformly over the whole images.
 
     Args:
-      images: A tensor of shape
-        (batch_size, height, width, channels)
-        (NHWC), (batch_size, channels, height, width)(NCHW).
+      images: A tensor of shape (batch_size, height, width, channels)(NHWC).
       mask_size: Specifies how big the zero mask that will be generated is that
         is applied to the images. The mask will be of size
         (mask_height x mask_width). Note: mask_size should be divisible by 2.
@@ -68,20 +56,13 @@ def random_cutout(
         the cutout mask applied to it.
       seed: A Python integer. Used in combination with `tf.random.set_seed` to
         create a reproducible sequence of tensors across multiple calls.
-      data_format: A string, one of `channels_last` (default) or `channels_first`.
-        The ordering of the dimensions in the inputs.
-        `channels_last` corresponds to inputs with shape
-        `(batch_size, ..., channels)` while `channels_first` corresponds to
-        inputs with shape `(batch_size, channels, ...)`.
     Returns:
       An image Tensor.
     Raises:
       InvalidArgumentError: if mask_size can't be divisible by 2.
     """
     batch_size = tf.shape(images)[0]
-    mask_size, data_format, image_height, image_width = _norm_params(
-        images, mask_size, data_format
-    )
+    mask_size, image_height, image_width = _norm_params(images, mask_size)
 
     cutout_center_height = tf.random.uniform(
         shape=[batch_size], minval=0, maxval=image_height, dtype=tf.int32, seed=seed
@@ -91,7 +72,7 @@ def random_cutout(
     )
 
     offset = tf.transpose([cutout_center_height, cutout_center_width], [1, 0])
-    return cutout(images, mask_size, offset, constant_values, data_format,)
+    return cutout(images, mask_size, offset, constant_values)
 
 
 def cutout(
@@ -99,29 +80,22 @@ def cutout(
     mask_size: TensorLike,
     offset: TensorLike = (0, 0),
     constant_values: Number = 0,
-    data_format: str = "channels_last",
 ) -> tf.Tensor:
     """Apply cutout (https://arxiv.org/abs/1708.04552) to images.
 
     This operation applies a (mask_height x mask_width) mask of zeros to
-    a location within `img` specified by the offset. The pixel values filled in will be of the
+    a location within `images` specified by the offset. The pixel values filled in will be of the
     value `replace`. The located where the mask will be applied is randomly
     chosen uniformly over the whole images.
 
     Args:
-      images: A tensor of shape (batch_size, height, width, channels)
-        (NHWC), (batch_size, channels, height, width)(NCHW).
+      images: A tensor of shape (batch_size, height, width, channels)(NHWC).
       mask_size: Specifies how big the zero mask that will be generated is that
         is applied to the images. The mask will be of size
         (mask_height x mask_width). Note: mask_size should be divisible by 2.
       offset: A tuple of (height, width) or (batch_size, 2)
       constant_values: What pixel value to fill in the images in the area that has
         the cutout mask applied to it.
-      data_format: A string, one of `channels_last` (default) or `channels_first`.
-        The ordering of the dimensions in the inputs.
-        `channels_last` corresponds to inputs with shape
-        `(batch_size, ..., channels)` while `channels_first` corresponds to
-        inputs with shape `(batch_size, channels, ...)`.
     Returns:
       An image Tensor.
     Raises:
@@ -130,9 +104,7 @@ def cutout(
     with tf.name_scope("cutout"):
         origin_shape = images.shape
         offset = tf.convert_to_tensor(offset)
-        mask_size, data_format, image_height, image_width = _norm_params(
-            images, mask_size, data_format
-        )
+        mask_size, image_height, image_width = _norm_params(images, mask_size)
         mask_size = mask_size // 2
 
         if tf.rank(offset) == 1:
@@ -165,12 +137,9 @@ def cutout(
             )
             masks = masks.write(i, mask)
 
-        if data_format == "channels_last":
-            mask_4d = tf.expand_dims(masks.stack(), -1)
-            mask = tf.tile(mask_4d, [1, 1, 1, tf.shape(images)[-1]])
-        else:
-            mask_4d = tf.expand_dims(masks.stack(), 1)
-            mask = tf.tile(mask_4d, [1, tf.shape(images)[1], 1, 1])
+        mask_4d = tf.expand_dims(masks.stack(), -1)
+        mask = tf.tile(mask_4d, [1, 1, 1, tf.shape(images)[-1]])
+
         images = tf.where(
             mask == 0,
             tf.ones_like(images, dtype=images.dtype) * constant_values,


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
@@ -48,13 +48,6 @@ def test_equalize_with_PIL():
     np.testing.assert_equal(color_ops.equalize(tf.constant(image)).numpy(), equalized)
 
 
-@pytest.mark.parametrize("shape", [(1, 5, 5), (3, 5, 5), (10, 3, 7, 7)])
-def test_equalize_channel_first(shape):
-    image = tf.ones(shape=shape, dtype=tf.uint8)
-    equalized = color_ops.equalize(image, "channels_first")
-    np.testing.assert_equal(equalized.numpy(), image.numpy())
-
-
 @pytest.mark.parametrize("dtype", _DTYPES)
 @pytest.mark.parametrize("shape", [(5, 5, 3), (10, 5, 5, 3)])
 def test_sharpness_dtype_shape(dtype, shape):
diff --git a/tensorflow_addons/image/tests/cutout_ops_test.py b/tensorflow_addons/image/tests/cutout_ops_test.py
--- a/tensorflow_addons/image/tests/cutout_ops_test.py
+++ b/tensorflow_addons/image/tests/cutout_ops_test.py
@@ -51,19 +51,6 @@ def test_batch_size():
     np.testing.assert_allclose(len(set(means)), 10)
 
 
-def test_channel_first():
-    test_image = tf.ones([10, 1, 40, 40], dtype=np.uint8)
-    cutout_area = tf.zeros([4, 4], dtype=np.uint8)
-    cutout_area = tf.pad(cutout_area, ((0, 36), (0, 36)), constant_values=1)
-    expect_image = tf.expand_dims(cutout_area, 0)
-    expect_image = tf.expand_dims(expect_image, 0)
-    expect_image = tf.tile(expect_image, [10, 1, 1, 1])
-    result_image = random_cutout(
-        test_image, 20, seed=1234, data_format="channels_first"
-    )
-    np.testing.assert_allclose(tf.shape(result_image), tf.shape(expect_image))
-
-
 @pytest.mark.usefixtures("maybe_run_functions_eagerly")
 def test_with_tf_function():
     test_image = tf.ones([1, 40, 40, 1], dtype=tf.uint8)