Add Augmenter utility

README.md

@@ -96,19 +96,20 @@ import keras_core as keras
 # Create a preprocessing pipeline with augmentations
 BATCH_SIZE = 16
 NUM_CLASSES = 3
-augmentations = [
-    keras_cv.layers.RandomFlip(),
-    keras_cv.layers.RandAugment(value_range=(0, 255)),
-    keras_cv.layers.CutMix(),
-]
+augmenter = keras_cv.layers.Augmenter(
+    [
+        keras_cv.layers.RandomFlip(),
+        keras_cv.layers.RandAugment(value_range=(0, 255)),
+        keras_cv.layers.CutMix(),
+    ],
+)
 
 def preprocess_data(images, labels, augment=False):
     labels = tf.one_hot(labels, NUM_CLASSES)
     inputs = {"images": images, "labels": labels}
     outputs = inputs
     if augment:
-        for augmentation in augmentations:
-            outputs = augmentation(outputs)
+        outputs = augmenter(outputs)
     return outputs['images'], outputs['labels']
 
 train_dataset, test_dataset = tfds.load(

keras_cv/layers/__init__.py

@@ -16,6 +16,7 @@
 from tensorflow.keras.layers import RandomHeight
 from tensorflow.keras.layers import RandomWidth
 
+from keras_cv.layers.augmenter import Augmenter
 from keras_cv.layers.feature_pyramid import FeaturePyramid
 from keras_cv.layers.fusedmbconv import FusedMBConvBlock
 from keras_cv.layers.mbconv import MBConvBlock

keras_cv/layers/augmenter.py

@@ -0,0 +1,43 @@
+# Copyright 2023 The KerasCV Authors
+#
+# 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
+#
+#     https://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.
+
+from keras_cv.backend import keras
+
+
+class Augmenter(keras.layers.Layer):
+    """Light-weight class to apply augmentations to data.
+
+    Args:
+        layers: A list of `keras.layers.Layers` to apply to the example
+
+    Examples:
+    from keras_cv import layers
+    images = np.ones((16, 256, 256, 3))
+    augmenter = layers.Augmenter(
+        [
+            layers.RandomFlip(),
+            layers.RandAugment(value_range=(0, 255)),
+            layers.CutMix(),
+        ]
+    )
+    augmented_images = augmenter(images)
+    """
+
+    def __init__(self, layers):
+        self.layers = layers
+
+    def __call__(self, inputs):
+        for layer in self.layers:
+            inputs = layer(inputs)
+        return inputs

keras_cv/layers/augmenter_test.py

@@ -0,0 +1,45 @@
+# Copyright 2023 The KerasCV Authors
+#
+# 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
+#
+#     https://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.
+
+from keras_cv import layers
+from keras_cv.backend import ops
+from keras_cv.tests.test_case import TestCase
+
+
+class AugmenterTest(TestCase):
+    def test_call(self):
+        images = ops.ones((2, 256, 256, 3))
+        augmenter = layers.Augmenter(
+            [
+                layers.RandomFlip(),
+                layers.RandAugment(value_range=(0, 255)),
+            ]
+        )
+        output = augmenter(images)
+        self.assertEquals(output.shape, images.shape)
+
+    def test_call_with_labels(self):
+        images = {
+            "labels": ops.ones((2,)),
+            "images": ops.ones((2, 256, 256, 3)),
+        }
+        augmenter = layers.Augmenter(
+            [
+                layers.RandomFlip(),
+                layers.RandAugment(value_range=(0, 255)),
+                layers.CutMix(),
+            ]
+        )
+        output = augmenter(images)
+        self.assertEquals(output["images"].shape, images["images"].shape)