How to flatten test image dataset and create a batch of tuple of (flattened image , labels)?

I'm working in Handwritten Math's symbol Classification using Federated Learning. I have preprocessed the image from keras.preprocessing.image.ImageDataGenerator and also obtained the labels of each images.

from keras.preprocessing.image import ImageDataGenerator

train_datagen = ImageDataGenerator(
        rescale=1./255,
        shear_range=0.2,
        zoom_range=0.2,
        horizontal_flip=True)
train_dataset = train_datagen.flow_from_directory(
        'train_test_data/train/',
        target_size=(45,45),
        batch_size=32,
        class_mode='categorical')

For obtaining labels:

import os
# make label list '!/exp87530.jpg'
def make_labels(train_dataset):
  labels = train_dataset.filenames
  label = []
  for l in labels:
    l = l.split(os.path.sep)[0]
    label.append(l)
  return label 

How could I make a tuple of flattened Image and label that needs to send to the clients ? As seen from tensorflow tutorial Building Your Own Federated Learning Algorithm

From tutorial:

import 
emnist_train, emnist_test = tff.simulation.datasets.emnist.load_data()

NUM_CLIENTS = 10
BATCH_SIZE = 20

def preprocess(dataset):

  def batch_format_fn(element):
    """Flatten a batch of EMNIST data and return a (features, label) tuple."""
    return (tf.reshape(element['pixels'], [-1, 784]), 
            tf.reshape(element['label'], [-1, 1]))

  return dataset.batch(BATCH_SIZE).map(batch_format_fn)

You can try something like this:

import tensorflow as tf

flowers = tf.keras.utils.get_file(
    'flower_photos',
    'https://storage.googleapis.com/download.tensorflow.org/example_images/flower_photos.tgz',
    untar=True)

img_gen = tf.keras.preprocessing.image.ImageDataGenerator(rescale=1./255,
        shear_range=0.2,
        zoom_range=0.2,
        horizontal_flip=True)


ds = tf.data.Dataset.from_generator(
    lambda: img_gen.flow_from_directory(flowers, target_size=(45,45), batch_size=10, shuffle=True),
    output_types=(tf.float32, tf.int32))

def preprocess(dataset):

  def batch_format_fn(x, y):
    """Flatten a batch of EMNIST data and return a (features, label) tuple."""
    return (tf.reshape(x, [-1, 45*45*3]), 
            tf.reshape(y, [-1, 5]))

  return dataset.map(batch_format_fn)

ds = preprocess(ds)
for x,y in ds.take(1):
  print(x.shape, y.shape)

Flattened batch of data, where 5 is the number of classes / different labels:

(10, 6075) (10, 5)