Repository: tensorflow/addons

Issue:
Cannot use random_cutout in keras training
**System information**
- OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Ubuntu 18.04
- TensorFlow version and how it was installed (source or binary): 2.1 from pip
- TensorFlow-Addons version and how it was installed (source or binary): 0.9.1 from pip
- Python version: Python 3.6.10 (default, Dec 19 2019, 23:04:32) [GCC 5.4.0 20160609] on linux
- Is GPU used? (yes/no): yes

**Describe the bug**

Cannot use the `tfa.image.random_cutout` is `tf.keras.Sequential.fit`;

**Code to reproduce the issue**

This MWE works:

 ```
import tensorflow as tf
import tensorflow_addons as tfa
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Conv2D, Dense, GlobalMaxPooling2D, Flatten

@tf.function
def preprocessing(image):
    return tf.image.convert_image_dtype(image, tf.float32)


X = tf.cast(tf.random.uniform(shape=(10, 224, 224, 3), minval=0, maxval=255), tf.uint8)
y = tf.random.uniform(shape=(10,))
dataset = tf.data.Dataset.from_tensor_slices((X, y)).map(lambda x, y: (preprocessing(x), y)).batch(3)
model = Sequential([Conv2D(10, (3, 3)), GlobalMaxPooling2D(), Flatten(), Dense(1)])
model.compile(loss="mse", optimizer="adam")
model.fit(dataset)
```

When I try to add the random_cutout, ie:

```
@tf.function
def preprocessing(image):
    return tf.image.convert_image_dtype(tf.squeeze(tfa.image.random_cutout(tf.expand_dims(image, 0), mask_size=(20, 20))), tf.float32)
```

it fails

**Other info / logs**

```
Traceback (most recent call last):
  File "test.py", line 16, in <module>
    model.fit(dataset)
  File "/home/clementw/Keras-FewShotLearning/venv/lib/python3.6/site-packages/tensorflow_core/python/keras/engine/training.py", line 819, in fit
    use_multiprocessing=use_multiprocessing)
  File "/home/clementw/Keras-FewShotLearning/venv/lib/python3.6/site-packages/tensorflow_core/python/keras/engine/training_v2.py", line 235, in fit
    use_multiprocessing=use_multiprocessing)
  File "/home/clementw/Keras-FewShotLearning/venv/lib/python3.6/site-packages/tensorflow_core/python/keras/engine/training_v2.py", line 593, in _process_training_inputs
    use_multiprocessing=use_multiprocessing)
  File "/home/clementw/Keras-FewShotLearning/venv/lib/python3.6/site-packages/tensorflow_core/python/keras/engine/training_v2.py", line 706, in _process_inputs
    use_multiprocessing=use_multiprocessing)
  File "/home/clementw/Keras-FewShotLearning/venv/lib/python3.6/site-packages/tensorflow_core/python/keras/engine/data_adapter.py", line 702, in __init__
    x = standardize_function(x)
  File "/home/clementw/Keras-FewShotLearning/venv/lib/python3.6/site-packages/tensorflow_core/python/keras/engine/training_v2.py", line 684, in standardize_function
    return dataset.map(map_fn, num_parallel_calls=dataset_ops.AUTOTUNE)
  File "/home/clementw/Keras-FewShotLearning/venv/lib/python3.6/site-packages/tensorflow_core/python/data/ops/dataset_ops.py", line 1591, in map
    self, map_func, num_parallel_calls, preserve_cardinality=True)
  File "/home/clementw/Keras-FewShotLearning/venv/lib/python3.6/site-packages/tensorflow_core/python/data/ops/dataset_ops.py", line 3926, in __init__
    use_legacy_function=use_legacy_function)
  File "/home/clementw/Keras-FewShotLearning/venv/lib/python3.6/site-packages/tensorflow_core/python/data/ops/dataset_ops.py", line 3147, in __init__
    self._function = wrapper_fn._get_concrete_function_internal()
  File "/home/clementw/Keras-FewShotLearning/venv/lib/python3.6/site-packages/tensorflow_core/python/eager/function.py", line 2395, in _get_concrete_function_internal
    *args, **kwargs)
  File "/home/clementw/Keras-FewShotLearning/venv/lib/python3.6/site-packages/tensorflow_core/python/eager/function.py", line 2389, in _get_concrete_function_internal_garbage_collected
    graph_function, _, _ = self._maybe_define_function(args, kwargs)
  File "/home/clementw/Keras-FewShotLearning/venv/lib/python3.6/site-packages/tensorflow_core/python/eager/function.py", line 2703, in _maybe_define_function
    graph_function = self._create_graph_function(args, kwargs)
  File "/home/clementw/Keras-FewShotLearning/venv/lib/python3.6/site-packages/tensorflow_core/python/eager/function.py", line 2593, in _create_graph_function
    capture_by_value=self._capture_by_value),
  File "/home/clementw/Keras-FewShotLearning/venv/lib/python3.6/site-packages/tensorflow_core/python/framework/func_graph.py", line 978, in func_graph_from_py_func
    func_outputs = python_func(*func_args, **func_kwargs)
  File "/home/clementw/Keras-FewShotLearning/venv/lib/python3.6/site-packages/tensorflow_core/python/data/ops/dataset_ops.py", line 3140, in wrapper_fn
    ret = _wrapper_helper(*args)
  File "/home/clementw/Keras-FewShotLearning/venv/lib/python3.6/site-packages/tensorflow_core/python/data/ops/dataset_ops.py", line 3082, in _wrapper_helper
    ret = autograph.tf_convert(func, ag_ctx)(*nested_args)
  File "/home/clementw/Keras-FewShotLearning/venv/lib/python3.6/site-packages/tensorflow_core/python/autograph/impl/api.py", line 237, in wrapper
    raise e.ag_error_metadata.to_exception(e)
ValueError: in converted code:

    /home/clementw/Keras-FewShotLearning/venv/lib/python3.6/site-packages/tensorflow_core/python/keras/engine/training_v2.py:677 map_fn
        batch_size=None)
    /home/clementw/Keras-FewShotLearning/venv/lib/python3.6/site-packages/tensorflow_core/python/keras/engine/training.py:2410 _standardize_tensors
        exception_prefix='input')
    /home/clementw/Keras-FewShotLearning/venv/lib/python3.6/site-packages/tensorflow_core/python/keras/engine/training_utils.py:529 standardize_input_data
        data = [standardize_single_array(x) for x in data]
    /home/clementw/Keras-FewShotLearning/venv/lib/python3.6/site-packages/tensorflow_core/python/keras/engine/training_utils.py:529 <listcomp>
        data = [standardize_single_array(x) for x in data]
    /home/clementw/Keras-FewShotLearning/venv/lib/python3.6/site-packages/tensorflow_core/python/keras/engine/training_utils.py:451 standardize_single_array
        if (x.shape is not None and len(x.shape) == 1 and
    /home/clementw/Keras-FewShotLearning/venv/lib/python3.6/site-packages/tensorflow_core/python/framework/tensor_shape.py:822 __len__
        raise ValueError("Cannot take the length of shape with unknown rank.")

    ValueError: Cannot take the length of shape with unknown rank.
```


Discussion:
I don't know if It Is caused  by the missing `img_utils.from_4D_image`.

There was a big thread closed yesterday about `Model.fit` with unknown ranks at https://github.com/tensorflow/tensorflow/issues/24520
I have updated tensorflow to 2.2 and the error is different:

```
ValueError: Input 0 of layer sequential is incompatible with the layer: its rank is undefined, but the layer requires a defined rank.
```
```
def preprocessing(image):
    cout_out = tfa.image.random_cutout(tf.expand_dims(image, 0), mask_size=(20, 20))
    output = tfa.image.utils.from_4D_image(cout_out,tfa.image.utils.get_ndims(image))
    return tf.dtypes.cast(output, tf.dtypes.float32)
```

Can you send a PR with something like https://github.com/tensorflow/addons/blob/master/tensorflow_addons/image/transform_ops.py#L117 for cut_out?
@ClementWalter if you need any feedback for the PR you could reach us also on https://gitter.im/tensorflow/sig-addons but I think it is quite trivial.

Resolving patch:
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
@@ -128,6 +128,7 @@ def cutout(
       InvalidArgumentError: if mask_size can't be divisible by 2.
     """
     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
@@ -175,4 +176,5 @@ def cutout(
             tf.ones_like(images, dtype=images.dtype) * constant_values,
             images,
         )
+        images.set_shape(origin_shape)
         return images


Test patch:
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
@@ -67,14 +67,8 @@ def test_channel_first():
 @pytest.mark.usefixtures("maybe_run_functions_eagerly")
 def test_with_tf_function():
     test_image = tf.ones([1, 40, 40, 1], dtype=tf.uint8)
-    result_image = tf.function(
-        random_cutout,
-        input_signature=[
-            tf.TensorSpec(shape=[None, 40, 40, 1], dtype=tf.uint8),
-            tf.TensorSpec(shape=[], dtype=tf.int32),
-        ],
-    )(test_image, 2)
+    result_image = tf.function(random_cutout)(test_image, 2)
     cutout_area = tf.zeros([4, 4], dtype=tf.uint8)
     cutout_area = tf.pad(cutout_area, ((0, 36), (0, 36)), constant_values=1)
     expect_image = to_4D_image(cutout_area)
-    np.testing.assert_allclose(tf.shape(result_image), tf.shape(expect_image))
+    np.testing.assert_equal(result_image.shape, expect_image.shape)