如何在MediaPipe Model Maker中指定特定的数据增强类型？

我正尝试使用MediaPipe Model Maker中的image_classifier训练自定义TFLite图像分类模型。目前仅能通过在HParams中设置布尔参数do_data_augmentation来开启随机增强（如裁剪、翻转等），代码示例如下：

spec = image_classifier.SupportedModels.MOBILENET_V2
hparams=image_classifier.HParams(epochs=100, export_dir="exported_model_2",do_data_augmentation=False)
options = image_classifier.ImageClassifierOptions(supported_model=spec, hparams=hparams)

model = image_classifier.ImageClassifier.create(
    train_data = train_data,
    validation_data = validation_data,
    options=options,
)

我希望指定数据增强的具体类型，仅保留翻转、曝光调整和模糊效果，请问是否可行？

可以通过替换默认的训练数据预处理流水线实现自定义增强，具体步骤如下：

1. 关闭默认自动增强：保持do_data_augmentation=False，避免和自定义增强逻辑冲突。
2. 编写自定义增强函数：用TensorFlow API实现你需要的翻转、曝光调整和模糊效果。
3. 将增强函数应用到训练数据集：修改train_data的dataset属性，替换为经过自定义增强处理后的数据集。

具体代码示例 #

1. 定义自定义增强函数

import tensorflow as tf

def custom_augment(image, label):
    # 随机水平翻转
    image = tf.image.random_flip_left_right(image)
    
    # 随机调整曝光（亮度），max_delta控制调整幅度，可按需修改
    image = tf.image.random_brightness(image, max_delta=0.2)
    
    # 高斯模糊实现（可自定义核大小和模糊程度）
    def apply_gaussian_blur(img):
        # 3x3高斯核
        kernel = tf.constant([[1, 2, 1],
                              [2, 4, 2],
                              [1, 2, 1]], dtype=tf.float32) / 16
        kernel = tf.expand_dims(tf.expand_dims(kernel, axis=-1), axis=-1)
        # 对单张图片添加batch维度后卷积，再移除batch维度
        return tf.nn.depthwise_conv2d(tf.expand_dims(img, 0), kernel, strides=[1,1,1,1], padding='SAME')[0]
    
    # 50%概率应用模糊（可调整触发概率）
    if tf.random.uniform(()) > 0.5:
        image = apply_gaussian_blur(image)
    
    # 确保像素值在合法范围（如果数据已归一化到0-1区间）
    image = tf.clip_by_value(image, 0.0, 1.0)
    
    return image, label

2. 应用自定义增强并训练模型

spec = image_classifier.SupportedModels.MOBILENET_V2
# 关闭默认增强
hparams = image_classifier.HParams(epochs=100, export_dir="exported_model_2", do_data_augmentation=False)
options = image_classifier.ImageClassifierOptions(supported_model=spec, hparams=hparams)

# 对训练数据应用自定义增强，启用多线程加速
train_data.dataset = train_data.dataset.map(custom_augment, num_parallel_calls=tf.data.AUTOTUNE)

# 创建并训练模型
model = image_classifier.ImageClassifier.create(
    train_data=train_data,
    validation_data=validation_data,
    options=options,
)

注意事项 #

• 像素值范围：如果训练数据未归一化（如0-255区间），需调整clip_by_value参数为0, 255，并确保亮度调整后的数值合法。
• 模糊效果：可通过修改高斯核大小（如5x5）或权重调整模糊程度，也可使用tf.keras.layers.GaussianNoise实现噪声类模糊。
• 增强概率：可为每个增强操作单独设置随机触发概率，让增强逻辑更灵活。

内容的提问来源于stack exchange，提问作者Michael Eilers Smith