Improve KerasCV Test routines

keras_cv/layers/preprocessing/auto_contrast_test.py

@@ -20,6 +20,28 @@
 
 
 class AutoContrastTest(TestCase):
+    def test_layer_basics(self):
+        img = np.array(
+            [[[1, 2, 3], [4, 5, 6]], [[7, 8, 9], [10, 11, 12]]],
+            dtype=np.float32,
+        )
+        img = np.expand_dims(img, axis=0)
+        expected_output = [
+            [
+                [[0.0, 0.0, 0.0], [85.0, 85.0, 85.0]],
+                [[170.0, 170.0, 170.0], [255.0, 255.0, 255.0]],
+            ]
+        ]
+        init_kwargs = {
+            "value_range": (0, 255),
+        }
+        self.run_preprocessing_layer_test(
+            cls=preprocessing.AutoContrast,
+            init_kwargs=init_kwargs,
+            input_data=img,
+            expected_output=expected_output,
+        )
+
     def test_constant_channels_dont_get_nanned(self):
         img = np.array([1, 1], dtype=np.float32)
         img = np.expand_dims(img, axis=-1)
@@ -60,16 +82,6 @@ def test_auto_contrast_different_values_per_channel(self):
         self.assertTrue(np.any(ops.convert_to_numpy(ys[0, ..., 0]) == 255.0))
         self.assertTrue(np.any(ops.convert_to_numpy(ys[0, ..., 1]) == 255.0))
 
-        self.assertAllClose(
-            ys,
-            [
-                [
-                    [[0.0, 0.0, 0.0], [85.0, 85.0, 85.0]],
-                    [[170.0, 170.0, 170.0], [255.0, 255.0, 255.0]],
-                ]
-            ],
-        )
-
     def test_auto_contrast_expands_value_range_uint8(self):
         img = np.array([0, 128], dtype=np.uint8)
         img = np.expand_dims(img, axis=-1)

keras_cv/layers/preprocessing/channel_shuffle_test.py

@@ -21,13 +21,6 @@
 
 
 class ChannelShuffleTest(TestCase):
-    def test_return_shapes(self):
-        xs = tf.ones((2, 512, 512, 3))
-
-        layer = ChannelShuffle(groups=3)
-        xs = layer(xs, training=True)
-        self.assertEqual(xs.shape, (2, 512, 512, 3))
-
     def test_channel_shuffle_call_results_one_channel(self):
         xs = tf.cast(
             tf.stack(
@@ -70,7 +63,21 @@ def test_non_square_image(self):
         self.assertTrue(np.any(ops.convert_to_numpy(xs[0]) == 2.0))
         self.assertTrue(np.any(ops.convert_to_numpy(xs[1]) == 1.0))
 
-    @pytest.mark.tf_only
+    def test_layer_basics(self):
+        input_data = tf.cast(
+            tf.ones((2, 512, 512, 3)),
+            dtype=tf.float32,
+        )
+        init_kwargs = {
+            "groups": 3,
+        }
+        self.run_preprocessing_layer_test(
+            cls=ChannelShuffle,
+            init_kwargs=init_kwargs,
+            input_data=input_data,
+            expected_output_shape=(2, 512, 512, 3),
+        )
+
     def test_in_tf_function(self):
         xs = tf.cast(
             tf.stack(

keras_cv/layers/preprocessing/cut_mix_test.py

@@ -23,6 +23,8 @@
 
 
 class CutMixTest(TestCase):
+    # Cannot do basic layer test since the layer cannot be tested on single
+    # image
     def test_return_shapes(self):
         xs = tf.ones((2, 512, 512, 3))
         # randomly sample labels

keras_cv/layers/preprocessing/equalization_test.py

@@ -23,13 +23,19 @@
 
 
 class EqualizationTest(TestCase):
-    def test_return_shapes(self):
-        xs = 255 * np.ones((2, 512, 512, 3), dtype=np.int32)
-        layer = Equalization(value_range=(0, 255))
-        xs = layer(xs)
-
-        self.assertEqual(xs.shape, (2, 512, 512, 3))
-        self.assertAllEqual(xs, 255 * np.ones((2, 512, 512, 3)))
+    def test_layer_basics(self):
+        input_data = 255 * np.ones((2, 512, 512, 3), dtype=np.int32)
+        init_kwargs = {
+            "value_range": (0, 255),
+        }
+        expected_output = 255 * np.ones((2, 512, 512, 3))
+        self.run_preprocessing_layer_test(
+            cls=Equalization,
+            init_kwargs=init_kwargs,
+            input_data=input_data,
+            expected_output=expected_output,
+            expected_output_shape=(2, 512, 512, 3),
+        )
 
     @pytest.mark.tf_keras_only
     def test_return_shapes_inside_model(self):

keras_cv/layers/preprocessing/fourier_mix_test.py

@@ -21,6 +21,8 @@
 
 
 class FourierMixTest(TestCase):
+    # Cannot do basic layer test since the layer cannot be tested on single
+    # image
     def test_return_shapes(self):
         xs = tf.ones((2, 512, 512, 3))
         # randomly sample labels

keras_cv/layers/preprocessing/grayscale_test.py

@@ -19,21 +19,17 @@
 
 
 class GrayscaleTest(TestCase):
-    def test_return_shapes(self):
-        xs = tf.ones((2, 52, 24, 3))
-
-        layer = preprocessing.Grayscale(
-            output_channels=1,
+    def test_layer_basics(self):
+        input_data = tf.ones((2, 52, 24, 3))
+        init_kwargs = {
+            "output_channels": 3,
+        }
+        self.run_preprocessing_layer_test(
+            cls=preprocessing.Grayscale,
+            init_kwargs=init_kwargs,
+            input_data=input_data,
+            expected_output_shape=(2, 52, 24, 3),
         )
-        xs1 = layer(xs, training=True)
-
-        layer = preprocessing.Grayscale(
-            output_channels=3,
-        )
-        xs2 = layer(xs, training=True)
-
-        self.assertEqual(xs1.shape, (2, 52, 24, 1))
-        self.assertEqual(xs2.shape, (2, 52, 24, 3))
 
     @pytest.mark.tf_only
     def test_in_tf_function(self):
@@ -66,41 +62,3 @@ def augment(x):
 
         self.assertEqual(xs1.shape, (2, 10, 10, 1))
         self.assertEqual(xs2.shape, (2, 10, 10, 3))
-
-    def test_non_square_image(self):
-        xs = tf.cast(
-            tf.stack([2 * tf.ones((52, 24, 3)), tf.ones((52, 24, 3))], axis=0),
-            tf.float32,
-        )
-
-        layer = preprocessing.Grayscale(
-            output_channels=1,
-        )
-        xs1 = layer(xs, training=True)
-
-        layer = preprocessing.Grayscale(
-            output_channels=3,
-        )
-        xs2 = layer(xs, training=True)
-
-        self.assertEqual(xs1.shape, (2, 52, 24, 1))
-        self.assertEqual(xs2.shape, (2, 52, 24, 3))
-
-    def test_in_single_image(self):
-        xs = tf.cast(
-            tf.ones((52, 24, 3)),
-            dtype=tf.float32,
-        )
-
-        layer = preprocessing.Grayscale(
-            output_channels=1,
-        )
-        xs1 = layer(xs, training=True)
-
-        layer = preprocessing.Grayscale(
-            output_channels=3,
-        )
-        xs2 = layer(xs, training=True)
-
-        self.assertEqual(xs1.shape, (52, 24, 1))
-        self.assertEqual(xs2.shape, (52, 24, 3))