Implementation of Cutout layer and RandomErasing layer (#47)

* initiate CutOut layer

* Implement CutOut layer

* docstring and rename

* initiate RandomErase layer

* Implement RandomErase layer

* refactor CutOut and RandomErase to inherit from same base class

* docstring

* remove unnecessary loop

* rename to noun instead of verb

* rename to noun instead of verb

* rename to noun instead of verb

* add get_config methods

* rename to noun instead of verb

* rename to noun instead of verb

* remove unnecessary check

* minor changes requested

* refactor fill_rectangle into a separate utils file

* import module instead of function

* refactor CutOut layer

* refactor RandomErasing layer

* add CutOut, RandomErasing to __init__.py

* add training argument to layer call

* rearrange imports

* remove labels argument from CutOut and RandomErasing

* add training argument to MixUp and CutMix

* update argument names for RandomErasing, and more detailed docstring.

* docstring fix

* docstring fix

* rearrange imports

* linting

* add license

* fix docstring

* fix docstring

* fix docstring

* refactor into single layer

* remove old layer

* refactor to make layer more extendable

* linting and additional tests

* black

* edit comment

* ceil height and width

* Update random_cutout to default to a 0 for the lower bound

* Remove unneeded if statement

Co-authored-by: Luke Wood <[email protected]>

Author: chjort

examples/layers/preprocessing/random_cutout_demo.py

@@ -0,0 +1,56 @@
+"""random_cutout_demo.py shows how to use the RandomCutout preprocessing layer.
+
+Operates on the oxford_flowers102 dataset.  In this script the flowers
+are loaded, then are passed through the preprocessing layers.
+Finally, they are shown using matplotlib.
+"""
+import matplotlib.pyplot as plt
+import tensorflow as tf
+import tensorflow_datasets as tfds
+
+from keras_cv.layers import preprocessing
+
+IMG_SIZE = (224, 224)
+BATCH_SIZE = 64
+
+
+def resize(image, label, num_classes=10):
+    image = tf.image.resize(image, IMG_SIZE)
+    label = tf.one_hot(label, num_classes)
+    return image, label
+
+
+def main():
+    data, ds_info = tfds.load(
+        "oxford_flowers102", with_info=True, as_supervised=True
+    )
+    train_ds = data["train"]
+
+    num_classes = ds_info.features["label"].num_classes
+
+    train_ds = (
+        train_ds.map(lambda x, y: resize(x, y, num_classes=num_classes))
+        .shuffle(10 * BATCH_SIZE)
+        .batch(BATCH_SIZE)
+    )
+    random_cutout = preprocessing.RandomCutout(
+        height_factor=(0.3, 0.9),
+        width_factor=64,
+        fill_mode="gaussian_noise",
+        rate=1.0,
+    )
+    train_ds = train_ds.map(
+        lambda x, y: (random_cutout(x), y), num_parallel_calls=tf.data.AUTOTUNE
+    )
+
+    for images, labels in train_ds.take(1):
+        plt.figure(figsize=(8, 8))
+        for i in range(9):
+            plt.subplot(3, 3, i + 1)
+            plt.imshow(images[i].numpy().astype("uint8"))
+            plt.axis("off")
+        plt.show()
+
+
+if __name__ == "__main__":
+    main()

keras_cv/layers/preprocessing/__init__.py

@@ -14,3 +14,4 @@
 
 from keras_cv.layers.preprocessing.cut_mix import CutMix
 from keras_cv.layers.preprocessing.mix_up import MixUp
+from keras_cv.layers.preprocessing.random_cutout import RandomCutout

keras_cv/layers/preprocessing/cut_mix.py

@@ -14,6 +14,9 @@
 import tensorflow as tf
 import tensorflow.keras.layers as layers
 from absl import logging
+from tensorflow.keras import backend
+
+from keras_cv.utils import fill_utils
 
 
 class CutMix(layers.Layer):
@@ -53,7 +56,7 @@ def _sample_from_beta(alpha, beta, shape):
         sample_beta = tf.random.gamma(shape, 1.0, beta=beta)
         return sample_alpha / (sample_alpha + sample_beta)
 
-    def call(self, images, labels):
+    def call(self, images, labels, training=True):
         """call method for the CutMix layer.
 
         Args:
@@ -66,6 +69,8 @@ def call(self, images, labels):
             labels: updated labels with both label smoothing and the cutmix updates
                 applied.
         """
+        if training is None:
+            training = backend.learning_phase()
 
         if tf.shape(images)[0] == 1:
             logging.warning(
@@ -74,9 +79,10 @@ def call(self, images, labels):
                 "expected.  Please call the layer with 2 or more samples."
             )
 
-        augment_cond = tf.less(
+        rate_cond = tf.less(
             tf.random.uniform(shape=[], minval=0.0, maxval=1.0), self.rate
         )
+        augment_cond = tf.logical_and(rate_cond, training)
         # pylint: disable=g-long-lambda
         cutmix_augment = lambda: self._update_labels(*self._cutmix(images, labels))
         no_augment = lambda: (images, self._smooth_labels(labels))
@@ -119,7 +125,7 @@ def _cutmix(self, images, labels):
         lambda_sample = tf.cast(lambda_sample, dtype=tf.float32)
 
         images = tf.map_fn(
-            lambda x: _fill_rectangle(*x),
+            lambda x: fill_utils.fill_rectangle(*x),
             (
                 images,
                 random_center_width,
@@ -148,48 +154,3 @@ def _smooth_labels(self, labels):
         off_value = label_smoothing / tf.cast(tf.shape(labels)[1], tf.float32)
         on_value = 1.0 - label_smoothing + off_value
         return labels * on_value + (1 - labels) * off_value
-
-
-def _fill_rectangle(
-    image, center_width, center_height, half_width, half_height, replace=None
-):
-    """Fill a rectangle in a given image using the value provided in replace.
-
-    Args:
-        image: the starting image to fill the rectangle on.
-        center_width: the X center of the rectangle to fill
-        center_height: the Y center of the rectangle to fill
-        half_width: 1/2 the width of the resulting rectangle
-        half_height: 1/2 the height of the resulting rectangle
-        replace: The value to fill the rectangle with.  Accepts a Tensor,
-            Constant, or None.
-    Returns:
-        image: the modified image with the chosen rectangle filled.
-    """
-    image_shape = tf.shape(image)
-    image_height = image_shape[0]
-    image_width = image_shape[1]
-
-    lower_pad = tf.maximum(0, center_height - half_height)
-    upper_pad = tf.maximum(0, image_height - center_height - half_height)
-    left_pad = tf.maximum(0, center_width - half_width)
-    right_pad = tf.maximum(0, image_width - center_width - half_width)
-
-    cutout_shape = [
-        image_height - (lower_pad + upper_pad),
-        image_width - (left_pad + right_pad),
-    ]
-    padding_dims = [[lower_pad, upper_pad], [left_pad, right_pad]]
-    mask = tf.pad(
-        tf.zeros(cutout_shape, dtype=image.dtype), padding_dims, constant_values=1
-    )
-    mask = tf.expand_dims(mask, -1)
-
-    if replace is None:
-        fill = tf.random.normal(image_shape, dtype=image.dtype)
-    elif isinstance(replace, tf.Tensor):
-        fill = replace
-    else:
-        fill = tf.ones_like(image, dtype=image.dtype) * replace
-    image = tf.where(tf.equal(mask, 0), fill, image)
-    return image

keras_cv/layers/preprocessing/mix_up.py

@@ -14,6 +14,7 @@
 import tensorflow as tf
 import tensorflow.keras.layers as layers
 from absl import logging
+from tensorflow.keras import backend
 
 
 class MixUp(layers.Layer):
@@ -53,7 +54,7 @@ def _sample_from_beta(alpha, beta, shape):
         sample_beta = tf.random.gamma(shape, 1.0, beta=beta)
         return sample_alpha / (sample_alpha + sample_beta)
 
-    def call(self, images, labels):
+    def call(self, images, labels, training=True):
         """call method for the MixUp layer.
 
         Args:
@@ -66,6 +67,8 @@ def call(self, images, labels):
             labels: updated labels with both label smoothing and the cutmix updates
                 applied.
         """
+        if training is None:
+            training = backend.learning_phase()
 
         if tf.shape(images)[0] == 1:
             logging.warning(
@@ -74,9 +77,10 @@ def call(self, images, labels):
                 "expected.  Please call the layer with 2 or more samples."
             )
 
-        augment_cond = tf.less(
+        rate_cond = tf.less(
             tf.random.uniform(shape=[], minval=0.0, maxval=1.0), self.rate
         )
+        augment_cond = tf.logical_and(rate_cond, training)
         # pylint: disable=g-long-lambda
         mixup_augment = lambda: self._update_labels(*self._mixup(images, labels))
         no_augment = lambda: (images, self._smooth_labels(labels))