Предварительная обработка DataAugmentation и MapDataset в федеративном обучении

Question

У меня есть следующий ImageDataGenerator:

      datagen = ImageDataGenerator(
    featurewise_center=False,  # set input mean to 0 over the dataset
    samplewise_center=False,  # set each sample mean to 0
    featurewise_std_normalization=False,  # divide inputs by std of the dataset
    samplewise_std_normalization=False,  # divide each input by its std
    zca_whitening=False,  # apply ZCA whitening
    rotation_range = 30,  # randomly rotate images in the range (degrees, 0 to 180)
    zoom_range = 0.2, # Randomly zoom image
    width_shift_range=0.1,  # randomly shift images horizontally (fraction of total width)
    height_shift_range=0.1,  # randomly shift images vertically (fraction of total height)
    horizontal_flip = True,  # randomly flip images
    vertical_flip=False)

И я хотел бы включить его в следующую функцию предварительной обработки:

      def preprocess(dataset):

 def batch_format_fn(element):
    a=tf.reshape(element['pixels'], [ -1, img_size, img_size , 1])
    """Flatten a batch `pixels` and return the features as an `OrderedDict`."""
    return collections.OrderedDict(
        x=a,
        y=tf.reshape(element['label'], [-1, 1]))
    d=dataset. \
    repeat(NUM_EPOCHS). \
    shuffle(SHUFFLE_BUFFER, seed=1). \
    batch(BATCH_SIZE). \
    map(batch_format_fn)
 return d

Как я могу включить datage.flow(..) в эту функцию в функции предварительной обработки ? Это проблема федеративного обучения