No public description

PiperOrigin-RevId: 609861906

official/vision/ops/augment.py

@@ -665,6 +665,7 @@ def _fill_rectangle_video(image,
   image_time = tf.shape(image)[0]
   image_height = tf.shape(image)[1]
   image_width = tf.shape(image)[2]
+  image_channels = tf.shape(image)[3]
 
   lower_pad = tf.maximum(0, center_height - half_height)
   upper_pad = tf.maximum(0, image_height - center_height - half_height)
@@ -681,7 +682,7 @@ def _fill_rectangle_video(image,
       padding_dims,
       constant_values=1)
   mask = tf.expand_dims(mask, -1)
-  mask = tf.tile(mask, [1, 1, 1, 3])
+  mask = tf.tile(mask, [1, 1, 1, image_channels])
 
   if replace is None:
     fill = tf.random.normal(tf.shape(image), dtype=image.dtype)

official/vision/ops/augment_test.py

@@ -479,12 +479,14 @@ def test_mixup_and_cutmix_smoothes_labels_with_videos(self):
     self.assertAllGreaterEqual(aug_labels, label_smoothing / num_classes -
                                1e4)  # With tolerance
 
-  def test_mixup_changes_video(self):
+  @parameterized.product(num_channels=[3, 4])
+  def test_mixup_changes_video(self, num_channels: int):
     batch_size = 12
     num_classes = 1000
     label_smoothing = 0.1
 
-    images = tf.random.normal((batch_size, 8, 224, 224, 3), dtype=tf.float32)
+    images = tf.random.normal(
+        (batch_size, 8, 224, 224, num_channels), dtype=tf.float32)
     labels = tf.range(batch_size)
     augmenter = augment.MixupAndCutmix(
         mixup_alpha=1., cutmix_alpha=0., num_classes=num_classes)
@@ -500,12 +502,14 @@ def test_mixup_changes_video(self):
                                1e4)  # With tolerance
     self.assertFalse(tf.math.reduce_all(images == aug_images))
 
-  def test_cutmix_changes_video(self):
+  @parameterized.product(num_channels=[3, 4])
+  def test_cutmix_changes_video(self, num_channels: int):
     batch_size = 12
     num_classes = 1000
     label_smoothing = 0.1
 
-    images = tf.random.normal((batch_size, 8, 224, 224, 3), dtype=tf.float32)
+    images = tf.random.normal(
+        (batch_size, 8, 224, 224, num_channels), dtype=tf.float32)
     labels = tf.range(batch_size)
     augmenter = augment.MixupAndCutmix(
         mixup_alpha=0., cutmix_alpha=1., num_classes=num_classes)