Introduction

ImageDataGenerator works well for simple single-input image classification, but it becomes awkward when a model has multiple inputs and an image-shaped target such as a segmentation mask. The practical solution is usually a custom generator or Sequence that coordinates augmentation across all inputs and targets, rather than trying to force one plain generator call to handle everything automatically.

Why the Simple Pattern Stops Working

The usual single-input pattern looks like:

datagen.flow(images, labels, batch_size=32)

That assumes:

• one image input
• one simple label target
• augmentation applied only to the input image

For multi-input or image-to-image tasks, you may need:

• two input images
• an auxiliary numeric or categorical input
• a target image that must stay spatially aligned with the input

That alignment requirement is the main reason a custom wrapper is needed.

Core Rule for Image Targets

If the target is an image mask or another spatial output, geometric augmentation must be synchronized between input and target.

For example, if the input image is rotated or flipped, the target mask must receive the exact same transform. Otherwise, the training pair becomes invalid.

That is why using one generator for the image and a separate unsynchronized generator for the mask is wrong.

A Practical Sequence Pattern

Here is a simple custom Sequence for two image inputs and one mask target.

1import math
2import numpy as np
3import tensorflow as tf
4
5from tensorflow.keras.preprocessing.image import ImageDataGenerator
6
7
8class MultiInputImageSequence(tf.keras.utils.Sequence):
9    def __init__(self, images_a, images_b, masks, batch_size=8):
10        self.images_a = images_a
11        self.images_b = images_b
12        self.masks = masks
13        self.batch_size = batch_size
14
15        self.image_gen = ImageDataGenerator(
16            rotation_range=10,
17            horizontal_flip=True,
18            rescale=1.0 / 255.0,
19        )
20
21        self.mask_gen = ImageDataGenerator(
22            rotation_range=10,
23            horizontal_flip=True,
24        )
25
26    def __len__(self):
27        return math.ceil(len(self.images_a) / self.batch_size)
28
29    def __getitem__(self, index):
30        start = index * self.batch_size
31        end = start + self.batch_size
32
33        batch_a = self.images_a[start:end]
34        batch_b = self.images_b[start:end]
35        batch_masks = self.masks[start:end]
36
37        seed = np.random.randint(0, 1_000_000)
38
39        gen_a = self.image_gen.flow(batch_a, batch_size=len(batch_a), shuffle=False, seed=seed)
40        gen_b = self.image_gen.flow(batch_b, batch_size=len(batch_b), shuffle=False, seed=seed)
41        gen_masks = self.mask_gen.flow(batch_masks, batch_size=len(batch_masks), shuffle=False, seed=seed)
42
43        augmented_a = next(gen_a)
44        augmented_b = next(gen_b)
45        augmented_masks = next(gen_masks)
46
47        return [augmented_a, augmented_b], augmented_masks

The shared seed is the important piece. It keeps geometric transforms aligned across the inputs and masks.

Example Model with Multiple Inputs

This model takes two images and predicts an image-like output.

1import tensorflow as tf
2
3input_a = tf.keras.Input(shape=(128, 128, 3), name="image_a")
4input_b = tf.keras.Input(shape=(128, 128, 3), name="image_b")
5
6x1 = tf.keras.layers.Conv2D(16, 3, padding="same", activation="relu")(input_a)
7x2 = tf.keras.layers.Conv2D(16, 3, padding="same", activation="relu")(input_b)
8
9x = tf.keras.layers.Concatenate()([x1, x2])
10x = tf.keras.layers.Conv2D(16, 3, padding="same", activation="relu")(x)
11output = tf.keras.layers.Conv2D(1, 1, activation="sigmoid")(x)
12
13model = tf.keras.Model(inputs=[input_a, input_b], outputs=output)
14model.compile(optimizer="adam", loss="binary_crossentropy")

Then train with the custom sequence:

# model.fit(sequence, epochs=5)

Important Mask Handling Note

Masks are not ordinary images. You usually do not want color jitter, brightness changes, or normalization intended for RGB inputs applied to the target mask. For masks, keep transformations limited to spatial transforms that preserve label meaning.

That is why the image generator and mask generator are not identical even when they share the same seed.

When tf.data Is the Better Choice

For complex pipelines, tf.data is often easier to scale and reason about than ImageDataGenerator, especially for:

• large datasets
• custom decoding
• mixed image and tabular inputs
• deterministic augmentation control

Still, if your codebase already uses ImageDataGenerator, wrapping it in a Sequence is a reasonable bridge solution.

Common Pitfalls

The most common mistake is augmenting input images and image targets independently, which destroys label alignment. Another is assuming ImageDataGenerator.flow can natively express any multi-input structure without a wrapper. Teams also often apply image normalization and color augmentation to segmentation masks, which corrupts target values. Finally, for complex multimodal pipelines, insisting on ImageDataGenerator alone can create brittle code where tf.data would be clearer.

Summary

• Plain ImageDataGenerator.flow is built for simpler single-input cases.
• For multiple inputs and image targets, use a custom generator or Sequence.
• Keep geometric augmentation synchronized across inputs and target images.
• Do not apply ordinary image color transforms blindly to masks.
• Consider tf.data when the pipeline becomes too complex for ImageDataGenerator.