如何将 imgaug 增强应用于 Tensorflow 2.0 中的 tf.dataDataset

Question

我有一个使用tf.data.Dataset图像和标签的输入管道的应用程序。现在我想使用增强功能，我正在尝试为此使用imgaug库。但是，我不知道该怎么做。我发现的所有示例都使用 KerasImageDataGenerator或Sequence.

在代码中，给定一个像这样的顺序增强器

  self.augmenter = iaa.Sequential([
        iaa.Fliplr(config.sometimes),
        iaa.Crop(percent=config.crop_percent),
        ...
        ], random_order=config.random_order)

我试图将该增强器应用于我的数据集中的批量图像，但没有成功。似乎我无法评估张量，因为我在 map 函数中运行我的增强。

def augment_dataset(self, dataset):
    dataset = dataset.map(self.augment_fn())
    return dataset

def augment_fn(self):
    def augment(images, labels):
        img_array = tf.make_ndarray(images)
        images = self.augmenter.augment_images(img_array) 
        return images, labels
    return augment

例如，如果我尝试使用 make_ndarray 我得到一个 AttributeError：'Tensor' object has no attribute 'tensor_shape'

这是因为 Dataset.map 没有使用 Eager 模式吗？。关于如何解决这个问题的任何想法？

更新 #1

我试过建议tf.numpy_function，如下

def augment_fn(self):
    def augment(images, labels):
        images = tf.numpy_function(self.augmenter.augment_images,
                                   [images],
                                   images.dtype)
        return images, labels
    return augment

但是，生成的图像具有未知的形状，这会导致稍后出现其他错误。如何保持图像的原始形状？在应用增强函数之前，我的一批图像有形状(batch_size, None, None, 1)，但之后形状是<unknown>

更新 #2

我通过首先找到图像的动态（真实）形状，然后重塑应用增强的结果来解决未知形状的问题

def augment_fn(self):
    def augment(images, labels):
        img_dtype = images.dtype
        img_shape = tf.shape(images)
        images = tf.numpy_function(self.augmenter.augment_images,
                                   [images],
                                   img_dtype)
        images = tf.reshape(images, shape = img_shape)
        return images, labels
    return augment

Answer 1

这是因为没有使用 eager 模式吗？我以为 TF2.0 中默认是 Eager 模式。关于如何解决这个问题有什么想法吗？

是的，数据集预处理不是在 eager 模式下执行的。我认为这是经过深思熟虑的，如果您认为数据集可以表示任意大（甚至无限）的数据流，那么这当然是有意义的。

假设你不可能/不切实际地将你正在做的增强转换为张量流操作（这将是第一选择！），那么你可以用来tf.numpy_function执行任意Python代码（这是现在已弃用的替代品tf.py_func）