def test_different_dtypes(dtype):
test_image = tf.ones([1, 40, 40, 1], dtype=dtype)
result_image = cutout(test_image, 4, [2, 2])
cutout_area = tf.zeros([4, 4], dtype=dtype)
cutout_area = tf.pad(cutout_area, ((0, 36), (0, 36)), constant_values=1)
expect_image = to_4D_image(cutout_area)
np.testing.assert_allclose(result_image, expect_image)
assert result_image.dtype == dtype
def test_invalid_mask_size():
with pytest.raises(tf.errors.InvalidArgumentError,
match="mask_size should be"):
x = tf.ones([10, 40, 40, 1], dtype=np.uint8)
cutout(x, 3, [2, 2]).numpy()
def test_keras_layer():
inputs = tf.keras.Input(shape=(40, 40, 1), dtype=tf.uint8)
outputs = tf.keras.layers.Lambda(lambda x: cutout(x, 4, [2, 2]))(inputs)
model = tf.keras.Model(inputs=inputs, outputs=outputs)
x = tf.ones([10, 40, 40, 1], dtype=np.uint8)
np.testing.assert_equal(cutout(x, 4, [2, 2]).numpy(), model(x))
