Introduction

Yes, ImageDataGenerator can be used for regression. The important difference from classification is that your target is a continuous value or vector, so you must feed raw numeric labels and choose augmentations that preserve the correctness of those labels.

The Core Rule for Regression

ImageDataGenerator does not care whether the target is a class index or a floating-point number. It simply yields image batches and matching target batches. That means regression works as long as:

• 'y is numeric'
• the model output shape matches the target shape
• the augmentation does not invalidate the label

For example, if you are predicting house price from an image, horizontal flip may be harmless. If you are predicting steering angle, a horizontal flip changes the correct target and must be handled carefully.

Example with Numpy Arrays

If your images are already loaded into memory, you can use flow directly:

1import numpy as np
2from tensorflow.keras import layers, models
3from tensorflow.keras.preprocessing.image import ImageDataGenerator
4
5X = np.random.rand(100, 128, 128, 3).astype("float32")
6y = np.random.rand(100, 1).astype("float32")
7
8datagen = ImageDataGenerator(
9    rescale=1.0 / 255.0,
10    rotation_range=10,
11    width_shift_range=0.1,
12    height_shift_range=0.1,
13    zoom_range=0.1,
14)
15
16train_gen = datagen.flow(X, y, batch_size=16, shuffle=True)
17
18model = models.Sequential([
19    layers.Input(shape=(128, 128, 3)),
20    layers.Conv2D(16, 3, activation="relu"),
21    layers.MaxPooling2D(),
22    layers.Conv2D(32, 3, activation="relu"),
23    layers.GlobalAveragePooling2D(),
24    layers.Dense(32, activation="relu"),
25    layers.Dense(1),
26])
27
28model.compile(optimizer="adam", loss="mse", metrics=["mae"])
29model.fit(train_gen, epochs=3)

The key regression-specific details are the final Dense(1) layer and the mse loss.

Example with DataFrames

When images live on disk and labels are in a CSV file, flow_from_dataframe is a common choice. For regression, set class_mode="raw".

1from tensorflow.keras.preprocessing.image import ImageDataGenerator
2
3datagen = ImageDataGenerator(rescale=1.0 / 255.0)
4
5train_gen = datagen.flow_from_dataframe(
6    dataframe=df,
7    directory="images",
8    x_col="filename",
9    y_col="target_value",
10    target_size=(128, 128),
11    batch_size=16,
12    class_mode="raw",
13    shuffle=True,
14)

class_mode="raw" is the crucial setting. Without it, Keras will try to interpret the target as a class label.

The same idea extends to vector-valued targets. If your regression output is, for example, two coordinates or several continuous measurements, the generator can still work as long as those numeric targets are emitted in the correct shape for the model output layer.

Be Careful with Label-Preserving Augmentation

Regression makes augmentation trickier than classification because not every image transform leaves the target unchanged.

Examples:

• brightness prediction: flips and crops may still be valid
• keypoint coordinates: transforms require target-coordinate updates
• steering angle: horizontal flip usually requires changing the sign of the label
• age estimation from faces: mild geometric transforms may be fine

So the right augmentation policy depends on what the numeric target actually means.

Modern Note

ImageDataGenerator still works, but many newer TensorFlow pipelines use tf.data plus preprocessing layers for augmentation. That approach is often easier to customize when regression targets also need transformation.

Still, if your use case is simple and labels remain valid under the augmentations, ImageDataGenerator is perfectly usable.

Common Pitfalls

• Forgetting class_mode="raw" when reading regression targets from a dataframe.
• Using a classification-style output layer or loss function for a regression problem.
• Applying augmentations that change the meaning of the target without updating the label.
• Assuming any image augmentation that works for classification is safe for regression.
• Normalizing images correctly but leaving target scaling unmanaged when the target range is large.

Summary

• 'ImageDataGenerator works for regression as long as the targets are fed as raw numeric values.'
• Use a regression output layer such as Dense(1) and a loss like mse or mae.
• With flow_from_dataframe, class_mode="raw" is usually the key setting.
• Only use augmentations that preserve or correctly transform the target value.
• For more complex target-aware augmentation, tf.data pipelines may be more flexible.