Repository: keras-team/keras-cv

Issue:
Support segmentation masks in `CutMix` layer
This should follow the same structure as segmentation mask augmentation in our other preprocessing layers.

Treating masks like 1-channel images for the purpose of applying CutMix should be a simple and effective approach for this.

Discussion:
cosmo3769: I will fix the conflicts as new code of adding support for PyTorch and JAX has been implemented in this.
ianstenbit: > I will fix the conflicts as new code of adding support for PyTorch and JAX has been implemented in this.

SG -- there are some merge conflicts, but keep in mind that right now, our preprocessing layers are TF-only (they are intended to run in the TF.data graph).

In the long term they'll support Jax+PyTorch but for now no need to do that for this PR
cosmo3769: Merge conflict resolved. Please check. Thank you.

Patch:
diff --git a/README.md b/README.md
index 34b2de7..cc84537 100644
--- a/README.md
+++ b/README.md
@@ -89,25 +89,26 @@ to take effect.
 
 ```python
 import tensorflow as tf
-from tensorflow import keras
 import keras_cv
 import tensorflow_datasets as tfds
+import keras_core as keras
 
 # Create a preprocessing pipeline with augmentations
 BATCH_SIZE = 16
 NUM_CLASSES = 3
-augmenter = keras.Sequential(
-    [
-        keras_cv.layers.RandomFlip(),
-        keras_cv.layers.RandAugment(value_range=(0, 255)),
-        keras_cv.layers.CutMix(),
-    ]
-)
+augmentations = [
+    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 = augmenter(inputs) if augment else inputs
+    outputs = inputs
+    if augment:
+        for augmentation in augmentations:
+            outputs = augmentation(outputs)
     return outputs['images'], outputs['labels']
 
 train_dataset, test_dataset = tfds.load(
diff --git a/examples/layers/preprocessing/segmentation/cut_mix_demo.py b/examples/layers/preprocessing/segmentation/cut_mix_demo.py
new file mode 100644
index 0000000..d9ba550
--- /dev/null
+++ b/examples/layers/preprocessing/segmentation/cut_mix_demo.py
@@ -0,0 +1,34 @@
+# 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.
+"""cut_mix_demo.py shows how to use the CutMix preprocessing layer.
+
+Uses the oxford iiit pet_dataset.  In this script the pets
+are loaded, then are passed through the preprocessing layers.
+Finally, they are shown using matplotlib.
+"""
+import demo_utils
+import tensorflow as tf
+
+from keras_cv.layers import preprocessing
+
+
+def main():
+    ds = demo_utils.load_oxford_iiit_pet_dataset()
+    cut_mix = preprocessing.CutMix(alpha=0.8)
+    ds = ds.map(cut_mix, num_parallel_calls=tf.data.AUTOTUNE)
+    demo_utils.visualize_dataset(ds)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/examples/layers/preprocessing/segmentation/demo_utils.py b/examples/layers/preprocessing/segmentation/demo_utils.py
index aa84b39..e7fc4dc 100644
--- a/examples/layers/preprocessing/segmentation/demo_utils.py
+++ b/examples/layers/preprocessing/segmentation/demo_utils.py
@@ -44,7 +44,6 @@ def to_dict(datapoint):
     input_mask = tf.one_hot(
         tf.squeeze(tf.cast(input_mask, tf.int32), axis=-1), depth=3
     )
-    input_image = tf.transpose(input_image, (2, 0, 1))
     return {"images": input_image, "segmentation_masks": input_mask}
 
 
@@ -78,5 +77,4 @@ def visualize_dataset(ds):
             samples["images"][0],
             samples["segmentation_masks"][0],
         )
-        sample_image = tf.transpose(sample_image, (1, 2, 0))
         display([sample_image, sample_mask])
diff --git a/keras_cv/layers/fusedmbconv.py b/keras_cv/layers/fusedmbconv.py
index 2d9dd86..5fff4ab 100644
--- a/keras_cv/layers/fusedmbconv.py
+++ b/keras_cv/layers/fusedmbconv.py
@@ -168,6 +168,13 @@ class FusedMBConvBlock(keras.layers.Layer):
             name=self.name + "project_bn",
         )
 
+        if self.survival_probability:
+            self.dropout = keras.layers.Dropout(
+                self.survival_probability,
+                noise_shape=(None, 1, 1, 1),
+                name=self.name + "drop",
+            )
+
     def build(self, input_shape):
         if self.name is None:
             self.name = keras.backend.get_uid("block0")
@@ -209,11 +216,7 @@ class FusedMBConvBlock(keras.layers.Layer):
         # Residual:
         if self.strides == 1 and self.input_filters == self.output_filters:
             if self.survival_probability:
-                x = keras.layers.Dropout(
-                    self.survival_probability,
-                    noise_shape=(None, 1, 1, 1),
-                    name=self.name + "drop",
-                )(x)
+                x = self.dropout(x)
             x = keras.layers.add([x, inputs], name=self.name + "add")
         return x
 
diff --git a/keras_cv/layers/mbconv.py b/keras_cv/layers/mbconv.py
index a94b2fd..bf6b104 100644
--- a/keras_cv/layers/mbconv.py
+++ b/keras_cv/layers/mbconv.py
@@ -175,6 +175,13 @@ class MBConvBlock(keras.layers.Layer):
             name=self.name + "project_bn",
         )
 
+        if self.survival_probability:
+            self.dropout = keras.layers.Dropout(
+                self.survival_probability,
+                noise_shape=(None, 1, 1, 1),
+                name=self.name + "drop",
+            )
+
     def build(self, input_shape):
         if self.name is None:
             self.name = keras.backend.get_uid("block0")
@@ -217,11 +224,7 @@ class MBConvBlock(keras.layers.Layer):
 
         if self.strides == 1 and self.input_filters == self.output_filters:
             if self.survival_probability:
-                x = keras.layers.Dropout(
-                    self.survival_probability,
-                    noise_shape=(None, 1, 1, 1),
-                    name=self.name + "drop",
-                )(x)
+                x = self.dropout(x)
             x = keras.layers.add([x, inputs], name=self.name + "add")
         return x
 
diff --git a/keras_cv/layers/preprocessing/cut_mix.py b/keras_cv/layers/preprocessing/cut_mix.py
index c521a7b..7b519eb 100644
--- a/keras_cv/layers/preprocessing/cut_mix.py
+++ b/keras_cv/layers/preprocessing/cut_mix.py
@@ -1,4 +1,4 @@
-# Copyright 2022 The KerasCV Authors
+# 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.
@@ -31,21 +31,18 @@ class CutMix(BaseImageAugmentationLayer):
             the smoothing values are sampled. Defaults to 1.0, which is a
             recommended value when training an imagenet1k classification model.
         seed: Integer. Used to create a random seed.
+
     References:
        - [CutMix paper]( https://arxiv.org/abs/1905.04899).
 
-    Sample usage:
-    ```python
-    (images, labels), _ = keras.datasets.cifar10.load_data()
-    labels = tf.one_hot(labels.squeeze(), 10)
-
-    cutmix = keras_cv.layers.preprocessing.cut_mix.CutMix(10)
-    output = cutmix({"images": images[:32], "labels": labels[:32]})
-    # output == {'images': updated_images, 'labels': updated_labels}
-    ```
     """
 
-    def __init__(self, alpha=1.0, seed=None, **kwargs):
+    def __init__(
+        self,
+        alpha=1.0,
+        seed=None,
+        **kwargs,
+    ):
         super().__init__(seed=seed, **kwargs)
         self.alpha = alpha
         self.seed = seed
@@ -63,16 +60,37 @@ class CutMix(BaseImageAugmentationLayer):
         self._validate_inputs(inputs)
         images = inputs.get("images", None)
         labels = inputs.get("labels", None)
+        segmentation_masks = inputs.get("segmentation_masks", None)
 
-        if images is None or labels is None:
-            raise ValueError(
-                "CutMix expects inputs in a dictionary with format "
-                '{"images": images, "labels": labels}.'
-                f"Got: inputs = {inputs}"
+        (
+            images,
+            lambda_sample,
+            permutation_order,
+            random_center_height,
+            random_center_width,
+            cut_width,
+            cut_height,
+        ) = self._cutmix(images)
+
+        if labels is not None:
+            labels = self._update_labels(
+                labels, lambda_sample, permutation_order
             )
-        images, labels = self._update_labels(*self._cutmix(images, labels))
+            inputs["labels"] = labels
+
+        if segmentation_masks is not None:
+            segmentation_masks = self._update_segmentation_masks(
+                segmentation_masks,
+                permutation_order,
+                random_center_height,
+                random_center_width,
+                cut_width,
+                cut_height,
+            )
+            inputs["segmentation_masks"] = segmentation_masks
+
         inputs["images"] = images
-        inputs["labels"] = labels
+
         return inputs
 
     def _augment(self, inputs):
@@ -82,7 +100,7 @@ class CutMix(BaseImageAugmentationLayer):
             "expected. Please call the layer with 2 or more samples."
         )
 
-    def _cutmix(self, images, labels):
+    def _cutmix(self, images):
         """Apply cutmix."""
         input_shape = tf.shape(images)
         batch_size, image_height, image_width = (
@@ -129,30 +147,60 @@ class CutMix(BaseImageAugmentationLayer):
 
         return (
             images,
-            tf.cast(labels, dtype=self.compute_dtype),
             lambda_sample,
             permutation_order,
+            random_center_height,
+            random_center_width,
+            cut_width,
+            cut_height,
         )
 
-    def _update_labels(self, images, labels, lambda_sample, permutation_order):
+    def _update_labels(self, labels, lambda_sample, permutation_order):
         cutout_labels = tf.gather(labels, permutation_order)
 
         lambda_sample = tf.reshape(lambda_sample, [-1, 1])
         labels = lambda_sample * labels + (1.0 - lambda_sample) * cutout_labels
-        return images, labels
+
+        return labels
+
+    def _update_segmentation_masks(
+        self,
+        segmentation_masks,
+        permutation_order,
+        random_center_height,
+        random_center_width,
+        cut_width,
+        cut_height,
+    ):
+        cutout_segmentation_masks = tf.gather(
+            segmentation_masks, permutation_order
+        )
+
+        segmentation_masks = fill_utils.fill_rectangle(
+            segmentation_masks,
+            random_center_width,
+            random_center_height,
+            cut_width,
+            cut_height,
+            cutout_segmentation_masks,
+        )
+
+        return segmentation_masks
 
     def _validate_inputs(self, inputs):
+        images = inputs.get("images", None)
         labels = inputs.get("labels", None)
-        if labels is None:
+        segmentation_masks = inputs.get("segmentation_masks", None)
+
+        if images is None or (labels is None and segmentation_masks is None):
             raise ValueError(
-                "CutMix expects 'labels' to be present in its inputs. "
-                "CutMix relies on both images an labels. "
-                "Please pass a dictionary with keys 'images' "
-                "containing the image Tensor, and 'labels' containing "
-                "the classification labels. "
-                "For example, `cut_mix({'images': images, 'labels': labels})`."
+                "CutMix expects inputs in a dictionary with format "
+                '{"images": images, "labels": labels}. or'
+                '{"images": images, "segmentation_masks": segmentation_masks}. '
+                f"Got: inputs = {inputs}."
             )
-        if not labels.dtype.is_floating:
+
+        if labels is not None and not labels.dtype.is_floating:
             raise ValueError(
                 f"CutMix received labels with type {labels.dtype}. "
                 "Labels must be of type float."
diff --git a/setup.py b/setup.py
index a652e99..4b39210 100644
--- a/setup.py
+++ b/setup.py
@@ -49,7 +49,13 @@ setup(
     author="Keras team",
     author_email="keras-cv@google.com",
     license="Apache License 2.0",
-    install_requires=["packaging", "absl-py", "regex", "tensorflow-datasets"],
+    install_requires=[
+        "packaging",
+        "absl-py",
+        "regex",
+        "tensorflow-datasets",
+        "keras-core",
+    ],
     python_requires=">=3.7",
     extras_require={
         "tests": [


Test patch:
diff --git a/keras_cv/layers/preprocessing/cut_mix_test.py b/keras_cv/layers/preprocessing/cut_mix_test.py
index 34bda0d..09f3983 100644
--- a/keras_cv/layers/preprocessing/cut_mix_test.py
+++ b/keras_cv/layers/preprocessing/cut_mix_test.py
@@ -1,4 +1,4 @@
-# Copyright 2022 The KerasCV Authors
+# 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.
@@ -23,18 +23,39 @@ class CutMixTest(TestCase):
     def test_return_shapes(self):
         xs = tf.ones((2, 512, 512, 3))
         # randomly sample labels
-        ys = tf.random.categorical(tf.math.log([[0.5, 0.5]]), 2)
-        ys = tf.squeeze(ys)
-        ys = tf.one_hot(ys, num_classes)
+        ys_labels = tf.random.categorical(tf.math.log([[0.5, 0.5]]), 2)
+        ys_labels = tf.squeeze(ys_labels)
+        ys_labels = tf.one_hot(ys_labels, num_classes)
+
+        # randomly sample segmentation mask
+        ys_segmentation_masks = tf.cast(
+            tf.stack(
+                [2 * tf.ones((512, 512)), tf.ones((512, 512))],
+                axis=0,
+            ),
+            tf.uint8,
+        )
+        ys_segmentation_masks = tf.one_hot(ys_segmentation_masks, 3)
 
         layer = CutMix(seed=1)
-        outputs = layer({"images": xs, "labels": ys})
-        xs, ys = outputs["images"], outputs["labels"]
+        outputs = layer(
+            {
+                "images": xs,
+                "labels": ys_labels,
+                "segmentation_masks": ys_segmentation_masks,
+            }
+        )
+        xs, ys_labels, ys_segmentation_masks = (
+            outputs["images"],
+            outputs["labels"],
+            outputs["segmentation_masks"],
+        )
 
         self.assertEqual(xs.shape, [2, 512, 512, 3])
-        self.assertEqual(ys.shape, [2, 10])
+        self.assertEqual(ys_labels.shape, [2, 10])
+        self.assertEqual(ys_segmentation_masks.shape, [2, 512, 512, 3])
 
-    def test_cut_mix_call_results(self):
+    def test_cut_mix_call_results_with_labels(self):
         xs = tf.cast(
             tf.stack(
                 [2 * tf.ones((4, 4, 3)), tf.ones((4, 4, 3))],
@@ -57,7 +78,7 @@ class CutMixTest(TestCase):
         self.assertNotAllClose(ys, 1.0)
         self.assertNotAllClose(ys, 0.0)
 
-    def test_cut_mix_call_results_one_channel(self):
+    def test_cut_mix_call_results_one_channel_with_labels(self):
         xs = tf.cast(
             tf.stack(
                 [2 * tf.ones((4, 4, 1)), tf.ones((4, 4, 1))],
@@ -80,6 +101,89 @@ class CutMixTest(TestCase):
         self.assertNotAllClose(ys, 1.0)
         self.assertNotAllClose(ys, 0.0)
 
+    def test_cut_mix_call_results_with_dense_encoded_segmentation_masks(self):
+        xs = tf.cast(
+            tf.stack(
+                [2 * tf.ones((4, 4, 3)), tf.ones((4, 4, 3))],
+                axis=0,
+            ),
+            tf.float32,
+        )
+
+        ys_segmentation_masks = tf.cast(
+            tf.stack(
+                [2 * tf.ones((4, 4, 1)), tf.ones((4, 4, 1))],
+                axis=0,
+            ),
+            tf.float32,
+        )
+
+        layer = CutMix(seed=1)
+        outputs = layer(
+            {"images": xs, "segmentation_masks": ys_segmentation_masks}
+        )
+        xs, ys_segmentation_masks = (
+            outputs["images"],
+            outputs["segmentation_masks"],
+        )
+
+        # At least some pixels should be replaced in the images
+        self.assertTrue(tf.math.reduce_any(xs[0] == 1.0))
+        self.assertTrue(tf.math.reduce_any(xs[0] == 2.0))
+        self.assertTrue(tf.math.reduce_any(xs[1] == 1.0))
+        self.assertTrue(tf.math.reduce_any(xs[1] == 2.0))
+        # At least some pixels should be replaced in the segmentation_masks
+        self.assertTrue(tf.math.reduce_any(ys_segmentation_masks[0] == 1.0))
+        self.assertTrue(tf.math.reduce_any(ys_segmentation_masks[0] == 2.0))
+        self.assertTrue(tf.math.reduce_any(ys_segmentation_masks[1] == 1.0))
+        self.assertTrue(tf.math.reduce_any(ys_segmentation_masks[1] == 2.0))
+
+    def test_cut_mix_call_results_with_one_hot_encoded_segmentation_masks(self):
+        xs = tf.cast(
+            tf.stack(
+                [2 * tf.ones((4, 4, 3)), tf.ones((4, 4, 3))],
+                axis=0,
+            ),
+            tf.float32,
+        )
+
+        ys_segmentation_masks = tf.cast(
+            tf.stack(
+                [2 * tf.ones((4, 4)), tf.ones((4, 4))],
+                axis=0,
+            ),
+            tf.uint8,
+        )
+        ys_segmentation_masks = tf.one_hot(ys_segmentation_masks, 3)
+
+        layer = CutMix(seed=1)
+        outputs = layer(
+            {"images": xs, "segmentation_masks": ys_segmentation_masks}
+        )
+        xs, ys_segmentation_masks = (
+            outputs["images"],
+            outputs["segmentation_masks"],
+        )
+
+        # At least some pixels should be replaced in the images
+        self.assertTrue(tf.math.reduce_any(xs[0] == 1.0))
+        self.assertTrue(tf.math.reduce_any(xs[0] == 2.0))
+        self.assertTrue(tf.math.reduce_any(xs[1] == 1.0))
+        self.assertTrue(tf.math.reduce_any(xs[1] == 2.0))
+        # At least some pixels should be replaced in the segmentation_masks
+        self.assertTrue(
+            tf.math.reduce_any(ys_segmentation_masks[0][:, :, 2] == 1.0)
+        )
+        self.assertTrue(
+            tf.math.reduce_any(ys_segmentation_masks[0][:, :, 2] == 0.0)
+        )
+        self.assertTrue(
+            tf.math.reduce_any(ys_segmentation_masks[1][:, :, 1] == 1.0)
+        )
+        self.assertTrue(
+            tf.math.reduce_any(ys_segmentation_masks[1][:, :, 1] == 0.0)
+        )
+
     def test_in_tf_function(self):
         xs = tf.cast(
             tf.stack(
@@ -117,13 +221,6 @@ class CutMixTest(TestCase):
         ):
             _ = layer(inputs)
 
-    def test_missing_labels(self):
-        xs = tf.ones((2, 512, 512, 3))
-        inputs = {"images": xs}
-        layer = CutMix()
-        with self.assertRaisesRegexp(ValueError, "CutMix expects 'labels'"):
-            _ = layer(inputs)
-
     def test_int_labels(self):
         xs = tf.ones((2, 512, 512, 3))
         ys = tf.one_hot(tf.constant([1, 0]), 2, dtype=tf.int32)
@@ -138,6 +235,6 @@ class CutMixTest(TestCase):
         xs = tf.ones((2, 512, 512, 3))
         layer = CutMix()
         with self.assertRaisesRegexp(
-            ValueError, "CutMix expects 'labels' to be present in its inputs"
+            ValueError, "CutMix expects inputs in a dictionary with format"
         ):
             _ = layer(xs)