► 开发者指南 / KerasCV / 使用 BaseImageAugmentationLayer 自定义图像增强

使用 BaseImageAugmentationLayer 自定义图像增强

作者： lukewood
创建日期： 2022/04/26
最后修改： 2023/11/29
描述： 使用 BaseImageAugmentationLayer 实现自定义数据增强。

概述

数据增强是训练任何强健的计算机视觉模型的一个重要组成部分。虽然 KerasCV 提供了大量现成的高质量数据增强技术，你可能仍然希望实现自己的自定义技术。 KerasCV 提供了一个有用的基类用于编写数据增强层： BaseImageAugmentationLayer。任何使用 BaseImageAugmentationLayer 构建的增强层将自动与 KerasCV 的 RandomAugmentationPipeline 类兼容。

本指南将向你展示如何使用 BaseImageAugmentationLayer 实现自己的自定义增强层。作为示例，我们将实现一个将所有图像着色为蓝色的层。

目前，KerasCV 的预处理层仅支持与 Keras 3 一起使用 TensorFlow 后端。

!pip install -q --upgrade keras-cv
!pip install -q --upgrade keras  # 升级到 Keras 3

import os

os.environ["KERAS_BACKEND"] = "tensorflow"

import keras
from keras import ops
from keras import layers
import keras_cv
import matplotlib.pyplot as plt

首先，让我们实现一些可视化和转换的辅助函数。

def imshow(img):
    img = img.astype(int)
    plt.axis("off")
    plt.imshow(img)
    plt.show()


def gallery_show(images):
    images = images.astype(int)
    for i in range(9):
        image = images[i]
        plt.subplot(3, 3, i + 1)
        plt.imshow(image.astype("uint8"))
        plt.axis("off")
    plt.show()


def transform_value_range(images, original_range, target_range):
    images = (images - original_range[0]) / (original_range[1] - original_range[0])
    scale_factor = target_range[1] - target_range[0]
    return (images * scale_factor) + target_range[0]


def parse_factor(param, min_value=0.0, max_value=1.0, seed=None):
    if isinstance(param, keras_cv.core.FactorSampler):
        return param
    if isinstance(param, float) or isinstance(param, int):
        param = (min_value, param)
    if param[0] == param[1]:
        return keras_cv.core.ConstantFactorSampler(param[0])
    return keras_cv.core.UniformFactorSampler(param[0], param[1], seed=seed)

BaseImageAugmentationLayer 介绍

图像增强应当以样本为基础进行操作，而不是以批次为基础。这是许多机器学习从业者在实现自定义技术时常犯的错误。 BaseImageAugmentation 提供了一组干净的抽象，使以样本为基础实现图像增强技术变得更简单。这通过允许最终用户重写 augment_image() 方法，并在内部执行自动向量化来实现。

大多数增强技术还必须从一个或多个随机分布中采样。 KerasCV 提供了一个抽象来使随机抽样的最终用户可配置：FactorSampler API。

最后，许多增强技术需要关于输入图像中像素值的一些信息。 KerasCV 提供了 value_range API 来简化对此的处理。

在我们的示例中，我们将使用 FactorSampler API、value_range API 和 BaseImageAugmentationLayer 来实现一个健壮、可配置且正确的 RandomBlueTint 层。

重写 `augment_image()`

让我们从最小开始：

class RandomBlueTint(keras_cv.layers.BaseImageAugmentationLayer):
    def augment_image(self, image, *args, transformation=None, **kwargs):
        # image 的形状是 (高度, 宽度, 通道)
        [*others, blue] = ops.unstack(image, axis=-1)
        blue = ops.clip(blue + 100, 0.0, 255.0)
        return ops.stack([*others, blue], axis=-1)

我们的层重写了 BaseImageAugmentationLayer.augment_image()。该方法用于增强传递给该层的图像。默认情况下，使用 BaseImageAugmentationLayer 会为你提供一些不错的功能：

支持未批处理的输入（HWC 张量）
支持批处理的输入（BHWC 张量）
对批处理输入的自动向量化（有关此的更多信息，请参见自动向量化性能）

让我们检查一下结果。首先，我们下载一张示例图像：

SIZE = (300, 300)
elephants = keras.utils.get_file(
    "african_elephant.jpg", "https://i.imgur.com/Bvro0YD.png"
)
elephants = keras.utils.load_img(elephants, target_size=SIZE)
elephants = keras.utils.img_to_array(elephants)
imshow(elephants)

从 https://i.imgur.com/Bvro0YD.png 下载数据
 4217496/4217496 ━━━━━━━━━━━━━━━━━━━━ 0s 0us/step

png

接下来，让我们进行增强并可视化结果：

layer = RandomBlueTint()
augmented = layer(elephants)
imshow(ops.convert_to_numpy(augmented))

png

看起来不错！我们也可以在批量输入上调用我们的层：

layer = RandomBlueTint()
augmented = layer(ops.expand_dims(elephants, axis=0))
imshow(ops.convert_to_numpy(augmented)[0])

png

使用 `FactorSampler` API 添加随机行为。

通常，图像增强技术在每次调用层的 __call__ 方法时不应该执行相同的操作。 KerasCV 提供了 FactorSampler API 以允许用户提供可配置的随机分布。

class RandomBlueTint(keras_cv.layers.BaseImageAugmentationLayer):
    """RandomBlueTint 随机地对图像应用蓝色色调。

    参数：
      factor：一个包含两个浮点数的元组，一个浮点数或
        `keras_cv.FactorSampler`。`factor` 控制图像的蓝移程度。`factor=0.0` 会使该层执行不操作，
        而 1.0 则完全使用退化结果。
        介于 0 和 1 之间的值将导致原始图像与完全蓝色图像之间的线性插值。
        值应在 `0.0` 和 `1.0` 之间。如果使用元组，则会为每个增强的图像在两个值之间采样一个 `factor`。
        如果使用单个浮点数，则在 `0.0` 和传递的浮点数之间采样一个值。为了确保值总是相同，
        请传递一个包含两个相同浮点数的元组：`(0.5, 0.5)`。
    """

    def __init__(self, factor, **kwargs):
        super().__init__(**kwargs)
        self.factor = parse_factor(factor)

    def augment_image(self, image, *args, transformation=None, **kwargs):
        [*others, blue] = ops.unstack(image, axis=-1)
        blue_shift = self.factor() * 255
        blue = ops.clip(blue + blue_shift, 0.0, 255.0)
        return ops.stack([*others, blue], axis=-1)

现在，我们可以配置我们的 RandomBlueTint 层的随机行为。我们可以给它一系列值进行抽样：

many_elephants = ops.repeat(ops.expand_dims(elephants, axis=0), 9, axis=0)
layer = RandomBlueTint(factor=0.5)
augmented = layer(many_elephants)
gallery_show(ops.convert_to_numpy(augmented))

png

每个图像的增强因素都是从范围(0, 0.5)的随机因素中抽样的。

我们还可以配置层以从正态分布中抽样：

many_elephants = ops.repeat(ops.expand_dims(elephants, axis=0), 9, axis=0)
factor = keras_cv.core.NormalFactorSampler(
    mean=0.3, stddev=0.1, min_value=0.0, max_value=1.0
)
layer = RandomBlueTint(factor=factor)
augmented = layer(many_elephants)
gallery_show(ops.convert_to_numpy(augmented))

png

如您所见，增强现在是从正态分布中抽样的。 FactorSamplers 有多种类型，包括 UniformFactorSampler、 NormalFactorSampler 和 ConstantFactorSampler。您也可以实现自己的。

重写 `get_random_transformation()`

现在，假设您的层影响预测目标：无论它们是边界框、分类标签还是回归目标。您的层需要了解在增强标签时对图像进行了哪些增强。幸运的是，BaseImageAugmentationLayer 是为此而设计的。

为了解决这个问题，BaseImageAugmentationLayer 具有可重写的 get_random_transformation() 方法，以及 augment_label()、 augment_target() 和 augment_bounding_boxes()。 augment_segmentation_map() 等其他功能将在将来添加。

让我们将此添加到我们的层中。

class RandomBlueTint(keras_cv.layers.BaseImageAugmentationLayer):
    """随机应用蓝色滤镜到图像。

    Args:
      factor: 一个包含两个浮点数的元组、一个单一的浮点数或一个
        `keras_cv.FactorSampler`。`factor` 控制图像的蓝色偏移程度。
        `factor=0.0` 使得该层执行无操作，值为 1.0 则完全使用退化结果。
        介于 0 和 1 之间的值则在原始图像与完全蓝色图像之间进行线性插值。
        值应介于 `0.0` 和 `1.0` 之间。如果使用元组，则为每个增强的图像
        在两个值之间随机采样一个 `factor`。如果使用单一浮点数，则将在 
        `0.0` 和传递的浮点数之间随机采样一个值。为了确保值始终相同，
        请传递一个包含两个相同浮点数的元组：(0.5, 0.5)。
    """

    def __init__(self, factor, **kwargs):
        super().__init__(**kwargs)
        self.factor = parse_factor(factor)

    def get_random_transformation(self, **kwargs):
        # kwargs 包含 {"images": image, "labels": label, 等等...}
        return self.factor() * 255

    def augment_image(self, image, transformation=None, **kwargs):
        [*others, blue] = ops.unstack(image, axis=-1)
        blue = ops.clip(blue + transformation, 0.0, 255.0)
        return ops.stack([*others, blue], axis=-1)

    def augment_label(self, label, transformation=None, **kwargs):
        # 如果需要，你可以以某种方式使用 transformation

        if transformation > 100:
            # 即，可能类别 2 对应于蓝色图像
            return 2.0

        return label

    def augment_bounding_boxes(self, bounding_boxes, transformation=None, **kwargs):
        # 你也可以对标签类型执行无操作增强，以支持它们在你的管道中。
        return bounding_boxes

To make use of these new methods, you will need to feed your inputs in with a dictionary maintaining a mapping from images to targets.

As of now, KerasCV supports the following label types:

labels via augment_label().
bounding_boxes via augment_bounding_boxes().

In order to use augmention layers alongside your prediction targets, you must package your inputs as follows:

labels = ops.array([[1, 0]])
inputs = {"images": ops.convert_to_tensor(elephants), "labels": labels}

Now if we call our layer on the inputs:

layer = RandomBlueTint(factor=(0.6, 0.6))
augmented = layer(inputs)
print(augmented["labels"])

2.0

输入和标签都被增强了。注意当 transformation 大于 100 时，标签被修改为包含 2.0，如上方层所述。

`value_range` 支持

想象一下，您在许多管道中使用新的增强层。一些管道的值在范围 [0, 255]，一些管道将其图像归一化到范围 [-1, 1]，还有一些使用的值范围是 [0, 1]。

如果用户用范围 [0, 1] 的图像调用您的层，输出将是无意义的！

layer = RandomBlueTint(factor=(0.1, 0.1))
elephants_0_1 = elephants / 255
print("最小值和最大值在增强之前:", elephants_0_1.min(), elephants_0_1.max())
augmented = layer(elephants_0_1)
print(
    "最小值和最大值在增强之后:",
    ops.convert_to_numpy(augmented).min(),
    ops.convert_to_numpy(augmented).max(),
)
imshow(ops.convert_to_numpy(augmented * 255).astype(int))

Clipping input data to the valid range for imshow with RGB data ([0..1] for floats or [0..255] for integers).

最小值和最大值在增强之前: 0.0 1.0
最小值和最大值在增强之后: 0.0 26.488235

png

请注意，这是一个非常弱的增强！因子仅设置为 0.1。

让我们用 KerasCV 的 value_range API 解决这个问题。

class RandomBlueTint(keras_cv.layers.BaseImageAugmentationLayer):
    """RandomBlueTint 随机给图像应用蓝色调。

    Args:
      value_range:  value_range: 一个包含两个元素的元组或列表。第一个值
        表示传入图像值的下限，第二个值表示上限。传入层的图像应具有
        在 `value_range` 之间的值。
      factor: 一个包含两个浮点数的元组、一个单一浮点数或
        `keras_cv.FactorSampler`。`factor` 控制图像蓝移的程度。`factor=0.0` 
        使此层执行无操作，而值为 1.0 则完全使用退化结果。
        在 0 和 1 之间的值会在原始图像和完全蓝色图像之间进行线性插值。
        值应在 `0.0` 和 `1.0` 之间。如果使用元组，则为每个增强的图像
        在这两个值之间采样一个 `factor`。如果使用单一浮点数，则在 `0.0` 和
        传入浮点数之间采样一个值。为了确保值始终相同，请传入两个相同浮点数的元组：
        `(0.5, 0.5)`。
    """

    def __init__(self, value_range, factor, **kwargs):
        super().__init__(**kwargs)
        self.value_range = value_range
        self.factor = parse_factor(factor)

    def get_random_transformation(self, **kwargs):
        # kwargs 包含 {"images": image, "labels": label, 等等...}
        return self.factor() * 255

    def augment_image(self, image, transformation=None, **kwargs):
        image = transform_value_range(image, self.value_range, (0, 255))
        [*others, blue] = ops.unstack(image, axis=-1)
        blue = ops.clip(blue + transformation, 0.0, 255.0)
        result = ops.stack([*others, blue], axis=-1)
        result = transform_value_range(result, (0, 255), self.value_range)
        return result

    def augment_label(self, label, transformation=None, **kwargs):
        # 如果需要，可以以某种方式使用 transformation

        if transformation > 100:
            # 即可能类 2 对应于蓝色图像
            return 2.0

        return label

    def augment_bounding_boxes(self, bounding_boxes, transformation=None, **kwargs):
        # 您也可以对标签类型执行无操作增强，以支持它们在
        # 您的管道中。
        return bounding_boxes


layer = RandomBlueTint(value_range=(0, 1), factor=(0.1, 0.1))
elephants_0_1 = elephants / 255
print("最小值和最大值在增强之前:", elephants_0_1.min(), elephants_0_1.max())
augmented = layer(elephants_0_1)
print(
    "最小值和最大值在增强之后:",
    ops.convert_to_numpy(augmented).min(),
    ops.convert_to_numpy(augmented).max(),
)
imshow(ops.convert_to_numpy(augmented * 255).astype(int))

增强前的最小值和最大值：0.0 1.0
增强后的最小值和最大值：0.0 1.0

png

现在我们的象只是稍微带有蓝色调。这是在使用 0.1 的情况下预期的行为。太好了！

现在用户可以配置该层以支持他们可能需要的任何值范围。请注意，只有与颜色信息交互的层应该使用值范围 API。许多增强技术，例如 RandomRotation 将不需要这个。

自动向量化性能

如果你在想：

在样本级别实现我的增强是否会带来性能影响？

你并不孤单！

幸运的是，我对自动向量化、手动向量化和未向量化实现的性能进行了广泛分析。在这个基准测试中，我使用自动向量化、无自动向量化和手动向量化实现了一个 RandomCutout 层。所有这些都在 @tf.function 注释中进行了基准测试。它们也在使用 jit_compile 参数时进行了基准测试。

以下图表显示了这个基准测试的结果：

自动向量化性能图表

主要结论是手动向量化和自动向量化之间的差异是微不足道的！

请注意，Eager 模式的性能将会有所不同。

常见问题

某些层无法自动向量化。一个例子是 GridMask。

如果你在调用你的层时收到错误，请尝试在构造函数中添加以下内容：

class UnVectorizable(keras_cv.layers.BaseImageAugmentationLayer):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        # 这将禁用 BaseImageAugmentationLayer 的自动向量化
        self.auto_vectorize = False

此外，请确保在 augment_* 方法中接受 **kwargs 以确保向前兼容性。KerasCV 将在未来增加其他标签类型，如果你不包含 **kwargs 参数，你的增强层将不会向前兼容。

结论和后续步骤

KerasCV 提供了一套标准的 API，以简化实现自己的数据增强技术的过程。这些包括 BaseImageAugmentationLayer、FactorSampler API 和 value_range API。

我们使用这些 API 实现了一个高度可配置的 RandomBlueTint 层。该层可以将输入作为独立图像、带有 "images" 和标签的字典、未分批的输入或分批的输入。输入可以在任何值范围内，并且用于采样色调值的随机分布是可由终端用户配置的。

作为后续练习，你可以：

使用 BaseImageAugmentationLayer 实现自己的数据增强技术
向 KerasCV 贡献一个增强层
浏览现有的 KerasCV 增强层

使用 BaseImageAugmentationLayer 自定义图像增强

◆ 概述

◆ BaseImageAugmentationLayer 介绍

◆ 重写 augment_image()

◆ 使用 FactorSampler API 添加随机行为。

◆ 重写 get_random_transformation()

◆ value_range 支持

◆ 自动向量化性能

◆ 常见问题

◆ 结论和后续步骤

使用 BaseImageAugmentationLayer 自定义图像增强

概述

BaseImageAugmentationLayer 介绍

重写 augment_image()

使用 FactorSampler API 添加随机行为。

重写 get_random_transformation()

value_range 支持

自动向量化性能

常见问题

结论和后续步骤

重写 `augment_image()`

使用 `FactorSampler` API 添加随机行为。

重写 `get_random_transformation()`

`value_range` 支持