tensorflow - 如何在 Tensorflow 2.0 中将 imgaug 增强应用于 tf.dataDataset

Question

我有一个带有输入管道的应用程序，它使用 tf.data.Dataset图像和标签。现在我想使用增强，我正在尝试使用 imgaug 为此目的的图书馆。但是，我不知道该怎么做。我发现的所有示例都使用 Keras ImageDataGenerator或 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 没有使用急切模式吗？关于如何解决这个问题的任何想法？

更新#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

我通过首先找到图像的动态(真实)形状然后 reshape 应用增强的结果来解决未知形状的问题

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

请转至TF Dataset documentation看看为什么在使用 tf.py_function 时需要返回图像的形状

def tf_random_rotate_image(image, label):
    im_shape = image.shape
    [image,] = tf.py_function(random_rotate_image, [image], [tf.float32])
    image.set_shape(im_shape)
    return image, label