使用Keras ImageDataGenerator可视化图像增强遇维度错误

问题场景 #

构建CNN模型训练德国交通标志图像（pickle格式），应用数据增强后，使用matplotlib和ImageDataGenerator显示增强图像时触发以下错误：

ValueError: ('Input data in NumpyArrayIterator should have rank 4. You passed an array with shape', (32, 32, 1))

错误原因 #

ImageDataGenerator.flow()要求输入数据是4维数组，格式为(样本数量, 图像高度, 图像宽度, 通道数)。而你传入的X_train_gray[i]是单张图像，维度为(32,32,1)，缺少了代表样本数的第一维，导致维度不匹配。

解决方法及修正代码 #

1. 为单张图像补充样本维度，将其从3维转为4维
2. 从数据生成器中获取增强后的图像批次
3. 调整图像维度以适配matplotlib的灰度图显示要求

修正后的可视化代码如下：

import numpy as np
import matplotlib.pyplot as plt

i = 100
# 方法1：用np.expand_dims增加样本维度，转为(1,32,32,1)
single_img = np.expand_dims(X_train_gray[i], axis=0)
# 方法2：用切片索引获取单样本，同样得到(1,32,32,1)
# single_img = X_train_gray[i:i+1]

# 传入生成器，获取增强后的图像批次
aug_iter = datagen.flow(single_img, batch_size=1)
# 获取第一个增强后的图像
aug_img = next(aug_iter)[0]

plt.figure(figsize=(10,8))
# 去掉通道维度，转为(32,32)以适配灰度图显示
plt.imshow(np.squeeze(aug_img), cmap='gray')
plt.show()

额外说明 #

如果需要生成多张增强图像，可以循环调用next(aug_iter)，并在子图中依次展示：

plt.figure(figsize=(15,10))
for j in range(6):
    aug_img = next(aug_iter)[0]
    plt.subplot(2,3,j+1)
    plt.imshow(np.squeeze(aug_img), cmap='gray')
    plt.axis('off')
plt.show()

内容的提问来源于stack exchange，提问作者Matthew_H